| import torch |
| import einops |
| import torchvision |
| import torch.nn as nn |
| from typing import List, Tuple |
|
|
|
|
| class MultiviewStack(nn.Module): |
| def __init__( |
| self, |
| encoders: List[nn.Module], |
| normalizations: List[Tuple[List, List]], |
| output_dim: int, |
| ): |
| super().__init__() |
| self.encoders = nn.ModuleList(encoders) |
| self.normalizations = [] |
| for mean, std in normalizations: |
| self.normalizations.append( |
| torchvision.transforms.Normalize(mean=mean, std=std) |
| ) |
|
|
| def forward(self, x): |
| orig_shape = x.shape |
| x = einops.rearrange(x, "... V C H W -> (...) V C H W") |
| outputs = [] |
| for i, encoder in enumerate(self.encoders): |
| this_view = x[:, i] |
| this_view = self.normalizations[i](this_view) |
| outputs.append(encoder(this_view)) |
| out = torch.stack(outputs, dim=-1) |
| out = out.reshape(*orig_shape[:-3], -1) |
| return out |
|
|
