def __call__()

in jcm/models/ncsnv2.py [0:0]

• 56 lines of code
• 2 McCabe index (conditional complexity)

    def __call__(self, x, labels, train=True):
        # config parsing
        config = self.config
        nf = config.model.nf
        act = get_act(config)
        normalizer = get_normalization(config)
        interpolation = config.model.interpolation

        if not config.data.centered:
            h = 2 * x - 1.0
        else:
            h = x

        h = conv3x3(h, nf, stride=1, bias=True)
        # ResNet backbone
        h = ResidualBlock(nf, resample=None, act=act, normalization=normalizer)(h)
        layer1 = ResidualBlock(nf, resample=None, act=act, normalization=normalizer)(h)
        h = ResidualBlock(2 * nf, resample="down", act=act, normalization=normalizer)(
            layer1
        )
        layer2 = ResidualBlock(
            2 * nf, resample=None, act=act, normalization=normalizer
        )(h)
        h = ResidualBlock(2 * nf, resample="down", act=act, normalization=normalizer)(
            layer2
        )
        layer3 = ResidualBlock(
            2 * nf, resample=None, act=act, normalization=normalizer
        )(h)
        h = ResidualBlock(
            4 * nf, resample="down", act=act, normalization=normalizer, dilation=2
        )(layer3)
        layer4 = ResidualBlock(
            4 * nf, resample=None, act=act, normalization=normalizer, dilation=2
        )(h)
        h = ResidualBlock(
            4 * nf, resample="down", act=act, normalization=normalizer, dilation=4
        )(layer4)
        layer5 = ResidualBlock(
            4 * nf, resample=None, act=act, normalization=normalizer, dilation=4
        )(h)
        # U-Net with RefineBlocks
        ref1 = RefineBlock(
            layer5.shape[1:3], 4 * nf, interpolation=interpolation, act=act, start=True
        )([layer5])
        ref2 = RefineBlock(
            layer4.shape[1:3], 2 * nf, interpolation=interpolation, act=act
        )([layer4, ref1])
        ref3 = RefineBlock(
            layer3.shape[1:3], 2 * nf, interpolation=interpolation, act=act
        )([layer3, ref2])
        ref4 = RefineBlock(layer2.shape[1:3], nf, interpolation=interpolation, act=act)(
            [layer2, ref3]
        )
        ref5 = RefineBlock(
            layer1.shape[1:3], nf, interpolation=interpolation, act=act, end=True
        )([layer1, ref4])

        h = normalizer()(ref5)
        h = act(h)
        h = conv3x3(h, x.shape[-1])

        return h