def __call__()

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

• 47 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, dilation=2
        )(layer2)
        layer3 = ResidualBlock(
            2 * nf, resample=None, act=act, normalization=normalizer, dilation=2
        )(h)
        h = ResidualBlock(
            2 * nf, resample="down", act=act, normalization=normalizer, dilation=4
        )(layer3)
        layer4 = ResidualBlock(
            2 * nf, resample=None, act=act, normalization=normalizer, dilation=4
        )(h)
        # U-Net with RefineBlocks
        ref1 = RefineBlock(
            layer4.shape[1:3], 2 * nf, act=act, interpolation=interpolation, start=True
        )([layer4])
        ref2 = RefineBlock(
            layer3.shape[1:3], 2 * nf, interpolation=interpolation, act=act
        )([layer3, ref1])
        ref3 = RefineBlock(
            layer2.shape[1:3], 2 * nf, interpolation=interpolation, act=act
        )([layer2, ref2])
        ref4 = RefineBlock(
            layer1.shape[1:3], nf, interpolation=interpolation, act=act, end=True
        )([layer1, ref3])

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

        return h