Class ResNetStack

tokenizer.py:303–419

ResNet encoder or decoder stack. Channel ratios and strides take the default order of from data/io-layer, to the middle of the model.

@si_module
class ResNetStack(nn.Module):
    """
    ResNet encoder or decoder stack. Channel ratios
    and strides take the default order of from 
    data/io-layer, to the middle of the model.
    """
    class Config:
        input_channels: int = 1
        output_channels: int = 1
        encode_channels: int = 32
        decode_channel_multiplier: int = 1
        latent_dim: int = None
        kernel_size: int = 7
        bias: bool = True
        channel_ratios: Tuple[int, ...] = (2, 4, 8, 16)
        strides: Tuple[int, ...] = (3, 4, 5, 5)
        mode: Literal['encoder', 'decoder'] = 'encoder'
        
    def __init__(self, c: Config):
        super().__init__()
        assert c.mode in ('encoder', 'decoder'), f"Mode ({c.mode}) is not supported!"
        
        self.mode = c.mode

        assert len(c.channel_ratios) == len(c.strides)
        channel_ratios = (1,) + c.channel_ratios
        strides = c.strides
        self.middle_channels = c.encode_channels * channel_ratios[-1]
        if c.mode == 'decoder':
            channel_ratios = tuple(reversed(channel_ratios))
            strides = tuple(reversed(strides))

        self.multiplier = c.decode_channel_multiplier if c.mode == 'decoder' else 1
        res_blocks = [ResNetBlock(
            c.encode_channels * channel_ratios[s_idx] * self.multiplier, 
            c.encode_channels * channel_ratios[s_idx+1] * self.multiplier, 
            stride,
            kernel_size=c.kernel_size,
            bias=c.bias,
            mode=c.mode,
        ) for s_idx, stride in enumerate(strides)]

        data_conv = CausalConv1d(
            in_channels=c.input_channels if c.mode == 'encoder' else c.encode_channels * self.multiplier, 
            out_channels=c.encode_channels if c.mode == 'encoder' else c.output_channels, 
            kernel_size=c.kernel_size, 
            stride=1, 
            bias=False,
        )

        if c.mode == 'encoder':
            self.res_stack = nn.Sequential(data_conv, *res_blocks)
        elif c.mode == 'decoder':
            self.res_stack = nn.Sequential(*res_blocks, data_conv)

        if c.latent_dim is not None:
            self.latent_proj = Conv1d1x1(self.middle_channels, c.latent_dim, bias=c.bias) if c.mode == 'encoder' else Conv1d1x1(c.latent_dim, self.middle_channels, bias=c.bias)
        if self.multiplier != 1: