Class ResnetBlock

src/encoding/blocks.py:108–161

class ResnetBlock(nn.Module):
    def __init__(self, in_channels, out_channels=None, dropout=0, up=False, num_groups=8, ks=3, input_norm=True, input_act=True):
        super().__init__()
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels
        self.up = up

        if input_norm and input_act:
            self.in_layers = nn.Sequential(
                nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True),
                SiLU(),
                nn.Conv2d(in_channels, out_channels, kernel_size=ks, stride=1, padding=(ks - 1)//2)
            )
        elif not input_norm:
            if input_act:
                self.in_layers = nn.Sequential(
                    SiLU(),
                    nn.Conv2d(in_channels, out_channels, kernel_size=ks, stride=1, padding=(ks - 1)//2)
                )
            else:
                self.in_layers = nn.Sequential(
                    nn.Conv2d(in_channels, out_channels, kernel_size=ks, stride=1, padding=(ks - 1)//2)
                )
        else:
            raise NotImplementedError

        self.out_layers = nn.Sequential(
            nn.GroupNorm(num_groups=num_groups, num_channels=out_channels, eps=1e-6, affine=True),
            SiLU(),
            nn.Dropout(p=dropout),
            zero_module(
                nn.Conv2d(out_channels, out_channels, kernel_size=ks, stride=1, padding=(ks - 1)//2)
            ),
        )

        if self.in_channels != self.out_channels:
            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        else:
            self.shortcut = nn.Identity()

    def forward(self, x):
        if self.up:
            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
            h = in_rest(x)
            h = F.interpolate(h, scale_factor=2, mode="bilinear", align_corners=False)
            x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
            h = in_conv(h)
        else:
            h = self.in_layers(x)
        
        h = self.out_layers(h)
        h = h + self.shortcut(x)
        return h


def compose_triplane_channelwise(feat_maps):