Apply projection of features.
| 60 | |
| 61 | |
| 62 | class ProjectionModule(nn.Module): |
| 63 | """Apply projection of features.""" |
| 64 | |
| 65 | def __init__(self, dims_in: list[int], dims_out: list[int]) -> None: |
| 66 | """Initialize projection module.""" |
| 67 | super().__init__() |
| 68 | if len(dims_in) != len(dims_out): |
| 69 | raise ValueError("Length of dims_in must be same as length of dims_out.") |
| 70 | self.convs = nn.ModuleList( |
| 71 | [nn.Conv2d(dim_in, dim_out, 1) for dim_in, dim_out in zip(dims_in, dims_out)] |
| 72 | ) |
| 73 | |
| 74 | def forward(self, encodings: list[torch.Tensor]) -> list[torch.Tensor]: |
| 75 | """Apply projection module.""" |
| 76 | if len(encodings) != len(self.convs): |
| 77 | raise ValueError("Number of encodings must be equal to number of projections.") |
| 78 | return [conv(encoding) for conv, encoding in zip(self.convs, encodings)] |
| 79 | |
| 80 | |
| 81 | class MonodepthFeatureEncoder(BaseEncoder): |