|
|
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 |
|
|
|
