(self, conf)
| 11 | |
| 12 | class EncoderBase(ABC, nn.Module): |
| 13 | def __init__(self, conf): |
| 14 | super().__init__() |
| 15 | self.conf = conf |
| 16 | self.t = conf.dataset.sequence_length |
| 17 | self.latent_channels = latent_channels = conf.model.latent_channels |
| 18 | self.separate_t_encoder = conf.model.separate_t_encoder |
| 19 | self.down_xyz = conf.model.down_x, conf.model.down_y, conf.model.down_z |
| 20 | |
| 21 | self.embedding = nn.Embedding(conf.dataset.num_classes, conf.model.latent_channels) |
| 22 | |
| 23 | self.voxel_encoder = VoxelEncoder(latent_channels, in_channels=latent_channels, down_xyz=self.down_xyz) |
| 24 | if self.separate_t_encoder: |
| 25 | t_in_channels = latent_channels + (self.conf.dataset.sequence_length if conf.model.one_hot_time else 0) |
| 26 | self.t_encoder = VoxelEncoder(latent_channels, in_channels=t_in_channels, down_xyz=self.down_xyz) |
| 27 | |
| 28 | self.norm = nn.InstanceNorm2d(latent_channels) |
| 29 | |
| 30 | down_ratios = [ |
| 31 | (conf.model.hex_down_x, conf.model.hex_down_y), |
| 32 | (conf.model.hex_down_x, conf.model.hex_down_z), |
| 33 | (conf.model.hex_down_y, conf.model.hex_down_z), |
| 34 | (conf.model.hex_down_t, conf.model.hex_down_x), |
| 35 | (conf.model.hex_down_t, conf.model.hex_down_y), |
| 36 | (conf.model.hex_down_t, conf.model.hex_down_z), |
| 37 | ] |
| 38 | self.downsamplers = nn.ModuleList( |
| 39 | [PlaneDownsampler(latent_channels, down_xy=ratio) for ratio in down_ratios] |
| 40 | ) |
| 41 | |
| 42 | def forward(self, x): |
| 43 | x = x.detach().clone() # B, T, X, Y, Z |
nothing calls this directly
no test coverage detected