neuralop.layers.embeddings.GridEmbedding2D
- class neuralop.layers.embeddings.GridEmbedding2D(in_channels: int, grid_boundaries=[[0, 1], [0, 1]])[source]
GridEmbedding2Dapplies a simple positional embedding as a regular 2D grid. Expects inputs of shape(batch, channels, d_1, d_2)- Parameters:
- in_channels
int number of channels in input. Fixed for output channel interface
- grid_boundaries
list, optional coordinate boundaries of input grid, by default [[0, 1], [0, 1]]
- in_channels
- Attributes:
- out_channels
Methods
forward(data[, batched])Define the computation performed at every call.
grid(spatial_dims, device, dtype)grid generates 2D grid needed for pos encoding and caches the grid associated with MRU resolution
- grid(spatial_dims, device, dtype)[source]
grid generates 2D grid needed for pos encoding and caches the grid associated with MRU resolution
- Parameters:
- spatial_dimstorch.size
sizes of spatial resolution
- deviceliteral ‘cpu’ or ‘cuda:*’
where to load data
- dtypestr
dtype to encode data
- Returns:
- torch.tensor
output grids to concatenate
- forward(data, batched=True)[source]
Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.