def __call__()

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

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

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

        return h