Class EncoderBase

dynamic_city/vae/encoder.py:12–69

class EncoderBase(ABC, nn.Module):
    def __init__(self, conf):
        super().__init__()
        self.conf = conf
        self.t = conf.dataset.sequence_length
        self.latent_channels = latent_channels = conf.model.latent_channels
        self.separate_t_encoder = conf.model.separate_t_encoder
        self.down_xyz = conf.model.down_x, conf.model.down_y, conf.model.down_z

        self.embedding = nn.Embedding(conf.dataset.num_classes, conf.model.latent_channels)

        self.voxel_encoder = VoxelEncoder(latent_channels, in_channels=latent_channels, down_xyz=self.down_xyz)
        if self.separate_t_encoder:
            t_in_channels = latent_channels + (self.conf.dataset.sequence_length if conf.model.one_hot_time else 0)
            self.t_encoder = VoxelEncoder(latent_channels, in_channels=t_in_channels, down_xyz=self.down_xyz)

        self.norm = nn.InstanceNorm2d(latent_channels)

        down_ratios = [
            (conf.model.hex_down_x, conf.model.hex_down_y),
            (conf.model.hex_down_x, conf.model.hex_down_z),
            (conf.model.hex_down_y, conf.model.hex_down_z),
            (conf.model.hex_down_t, conf.model.hex_down_x),
            (conf.model.hex_down_t, conf.model.hex_down_y),
            (conf.model.hex_down_t, conf.model.hex_down_z),
        ]
        self.downsamplers = nn.ModuleList(
            [PlaneDownsampler(latent_channels, down_xy=ratio) for ratio in down_ratios]
        )

    def forward(self, x):
        x = x.detach().clone()  # B, T, X, Y, Z
        x = self.embedding(x)  # B, T, X, Y, Z, C
        t_vox = None
        if self.separate_t_encoder:
            t_vox = self.vox_convs(x, self.t_encoder, one_hot=self.conf.model.one_hot_time)
        x = self.vox_convs(x, self.voxel_encoder)
        x = self.vox_to_planes(x, t_vox)
        for i, downsampler in enumerate(self.downsamplers):
            x[i] = self.downsamplers[i](x[i])
        return x

    def vox_convs(self, x, encoder, one_hot=False):
        B, T, X, Y, Z, C = x.shape
        if one_hot:
            x = rearrange(x, 'b t x y z c -> b t c x y z')
            one_hot_tensor = torch.eye(T, device=x.device)[None, :, :, None, None, None].expand(B, T, T, X, Y, Z)
            x = torch.cat([x, one_hot_tensor], dim=2)
            x = rearrange(x, 'b t c x y z -> (b t) c x y z')
        else:
            x = rearrange(x, 'b t x y z c -> (b t) c x y z')
        x = encoder(x)
        x = rearrange(x, '(b t) c x y z -> b c t x y z', b=B, t=T)
        return x

    @abstractmethod
    def vox_to_planes(self, x, t_vox=None):
        pass