Add option to output raw patch embeddings

[yellowcap]

This PR adds an option to the Lightning CLI to output patch level embeddings. I.e. one embedding per patch in each chip. The band group dimension is reduced and, so the patch embeddings are averages over the band groups.

This adds two args to the CLI output_patch_embeddings, and shuffle. Because for the patch embeddings we need to ensure that shuffle is off, while it is on by default. See also #123.

Updated the documentation to explain these changes.

This PR Closes #130

[srmsoumya]

The order of the tensor matters, in this case the ouput from the encoder i.e embeddings_raw is of shape batch x (group x num_spatial_patches) x embedding_dims.

So, the einops operation should be, b (g l) d -> b g l d - you can check this notebook for reference: https://github.com/Clay-foundation/model/blob/docs/model/docs/clay-v0-visualization.ipynb

[yellowcap]

Done, but please review again to make sure I got it right this time. Related to this, do you think this is a good way to unravel the patch embeddings?

[weiji14]

Mentioned this below at #133 (comment), but can we just keep the unravelled shape (B, 256, 768), or (B, 1536, 768)? Less work for the downstream user since they won't need to figure out how to unravel the tensor.

[yellowcap]

Yes absolutely. Don't know why I thought it had to be a 1d array...

[weiji14]

That said, do we want to store a (256, 768) embedding tensor in a single row, or split it into 256 rows of 768-length embeddings? I'm reading the original thread at #127, and it sounds like we need to enable searching on the 32x32 patch level instead of 256x256 chips, which might mean storing a row for each patch?

We can store 2D arrays in GeoParquet, but I'm not sure if vector databases allow indexing 2D arrays, or if we need to make it 1D. Note that this would increase the embedding file size significantly (x256), and the vector database indexing will be much slower. But if these raw embeddings are only meant to be generated on an ad-hoc basis for small locations, and not for similarity search applications, it should be fine.

[yellowcap]

I thought about that too, but that would also imply that we have to track the MGRS tile and the source path 256x more. That seems ineficient. My thinking was that the splitting into rows can happen at the moment on ingestion to a vector search DB or when analysing if required. For "local manual inspection" situations, having the patches in a single array is useful too. So I vote to keep it at one row in the gpd and let the separation be a downstream issue.

[weiji14]

Dangling sentence, representing what? Also, should be size 16*16, 768 no?

[yellowcap]

Updated wording to adopt to the embedding levels input.

[weiji14]

If we decide to use the unsqueezed shape.

Suggested change

	requested, the shape is (B, 16 * 16 * 768), one embedding per patch.
	requested, the shape is (B, 6*16*16, 768), one embedding per patch.

[weiji14]

Why are we squeezing the embeddings to a size like (B, 256*768)?? Can't we just keep the raw embedding shape of (B, 1538, 768), where 1538 is 6 (band groups) * 16 * 16 (patch size)? It will be a lot harder to unsqueeze a 196608 length tensor back to 6*16*16 later.

[yellowcap]

Unsqueezed would undoubtedly be better, but I assumed that we can only store one-dimensional arrays in a field. I guess from your questions that is not the case. Can we store multidimensional arrays here? If yes, lets store the unsqueezed one for sure!

[weiji14]

Yes, GeoParquet/Arrow technically supports FixedShapeTensorArray which can be multi-dimensional.

[yellowcap]

I just found out that this will require a rewrite without pandas, because pandas can not handle this and implicitly flattens the FixedShapeTensorArray, see example below. I ran across this when running tests and not getting the expected multidimensional arrays. This defintively complicates things in terms of keeping the structure.

import numpy as np
import pandas as pd
import pyarrow as pa

In [15]: array = np.arange(8).reshape((2,2,2))
    ...: 
    ...: arrow = pa.FixedShapeTensorArray.from_numpy_ndarray(array)
    ...: 
    ...: df = pd.DataFrame(data={"arrow": arrow})
    ...: 
    ...: df
    ...: 
Out[15]: 
          arrow
0  [0, 1, 2, 3]
1  [4, 5, 6, 7]

[yellowcap]

Ok so based on the above, the arrays are outputted in a flat structure for now. I updated the documentation with a sentence about where the structure comes from and why the arrays are flat.

https://github.com/Clay-foundation/model/pull/133/files#diff-2f39d0012d2540cdab1ac1ff882a3875d829b85f0acee68fe102e6b8a4f2c6c0R60

This is not ideal but keeping the multidimensional arrays now would require refactoring of how we construct the geoparquet files.

[weiji14]

Hmm, have you tried inserting a FixedShapeTensorArray into a row cell directly? This seems to work:

import numpy as np
import pandas as pd
import pyarrow as pa

array0 = np.arange(0, 8).reshape((2, 2, 2))
array1 = np.arange(9, 17).reshape((2, 2, 2))

arrow0 = pa.FixedShapeTensorArray.from_numpy_ndarray(array0)
arrow1 = pa.FixedShapeTensorArray.from_numpy_ndarray(array1)

df = pd.DataFrame(data={"embeddings": [arrow0, arrow1]})
print(df)
#                             embeddings
# 0         ([0, 1, 2, 3], [4, 5, 6, 7])
# 1  ([9, 10, 11, 12], [13, 14, 15, 16])

[yellowcap]

The arrow part still flattens the arrays that are input. For this to work, we would have to do the manual deconstruction up to 3 dimensions deep. Not sure that this will make things more user friendly. With this approach, the user would have to reconstruct a list of a list of a list of arrays in the "group" case.

[srmsoumya]

Like the idea of storing embeddings at multiple levels.

[srmsoumya]

Is bf16-mixed precision giving any issue while inferencing?

[yellowcap]

It is for the Nvidia Geforce RTX on my local machine, which does not support this type apparently.