def __init__(self, in_channels, hidden_channels, actnorm_scale,
                 flow_permutation, flow_coupling, LU_decomposed):
        super().__init__()
        self.flow_coupling = flow_coupling

        self.actnorm = ActNorm2d(in_channels, actnorm_scale)

        # 2. permute
        if flow_permutation == "invconv":
            self.invconv = InvertibleConv1x1(in_channels,
                                             LU_decomposed=LU_decomposed)
            self.flow_permutation = \
                lambda z, logdet, rev: self.invconv(z, logdet, rev)
        elif flow_permutation == "shuffle":
            self.shuffle = Permute2d(in_channels, shuffle=True)
            self.flow_permutation = \
                lambda z, logdet, rev: (self.shuffle(z, rev), logdet)
        else:
            self.reverse = Permute2d(in_channels, shuffle=False)
            self.flow_permutation = \
                lambda z, logdet, rev: (self.reverse(z, rev), logdet)

        # 3. coupling
        if flow_coupling == "additive":
            self.block = get_block(in_channels // 2, in_channels // 2,
                                   hidden_channels)
        elif flow_coupling == "affine":
            self.block = get_block(in_channels // 2, in_channels,
                                   hidden_channels)
Пример #2
0
    def __init__(self, in_channels, hidden_channels, actnorm_scale,
                 flow_permutation, flow_coupling, LU_decomposed,
                 flow_embed_dim, level):
        super().__init__()
        self.flow_coupling = flow_coupling

        self.actnorm = ActNorm2d(in_channels, actnorm_scale)
        self.actnorm_embed = ActNorm2d(flow_embed_dim * (2**level),
                                       actnorm_scale)

        if flow_coupling == "additive":
            self.conv_proj = Conv2dZeros(hidden_channels,
                                         in_channels // 2,
                                         kernel_size=(3, 3))
        elif flow_coupling == "affine":
            self.conv_proj = Conv2dZeros(hidden_channels,
                                         in_channels,
                                         kernel_size=(3, 3))

        # 2. permute
        if flow_permutation == "invconv":
            self.invconv = InvertibleConv1x1(in_channels,
                                             LU_decomposed=LU_decomposed)
            self.flow_permutation = \
                lambda z, logdet, rev: self.invconv(z, logdet, rev)
        elif flow_permutation == "shuffle":
            self.shuffle = Permute2d(in_channels, shuffle=True)
            self.flow_permutation = \
                lambda z, logdet, rev: (self.shuffle(z, rev), logdet)
        else:
            self.reverse = Permute2d(in_channels, shuffle=False)
            self.flow_permutation = \
                lambda z, logdet, rev: (self.reverse(z, rev), logdet)

        self.multgate = nn.Parameter(torch.zeros((6, hidden_channels, 1, 1)))
Пример #3
0
    def __init__(self,
                 in_channels,
                 hidden_channels,
                 actnorm_scale,
                 flow_permutation,
                 flow_coupling,
                 LU_decomposed,
                 sn,
                 affine_eps,
                 no_actnorm,
                 affine_scale_eps=0,
                 actnorm_max_scale=None,
                 no_conv_actnorm=False,
                 affine_max_scale=0,
                 actnorm_eps=0):
        super().__init__()
        self.flow_coupling = flow_coupling
        self.affine_eps = affine_eps
        self.max_scale = affine_max_scale
        assert not (flow_coupling == 'gaffine' and affine_max_scale == 0)
        self.no_actnorm = no_actnorm
        self.affine_scale_eps = affine_scale_eps
        if not self.no_actnorm:
            self.actnorm = ActNorm2d(in_channels,
                                     actnorm_scale,
                                     max_scale=actnorm_max_scale,
                                     actnorm_eps=actnorm_eps)

        # 2. permute
        if flow_permutation == "invconv":
            self.invconv = InvertibleConv1x1(in_channels,
                                             LU_decomposed=LU_decomposed)
            self.flow_permutation = \
                lambda z, logdet, rev: self.invconv(z, logdet, rev)
        elif flow_permutation == "shuffle":
            self.shuffle = Permute2d(in_channels, shuffle=True)
            self.flow_permutation = \
                lambda z, logdet, rev: (self.shuffle(z, rev), logdet)
        else:
            self.reverse = Permute2d(in_channels, shuffle=False)

        # 3. coupling
        if flow_coupling == "additive":
            self.block = get_block(in_channels // 2, in_channels // 2,
                                   hidden_channels, sn, no_conv_actnorm)
        elif flow_coupling in ["affine", "naffine", "gaffine"]:
            self.block = get_block(in_channels // 2, in_channels,
                                   hidden_channels, sn, no_conv_actnorm)
Пример #4
0
    def __init__(self, in_channels, hidden_channels, actnorm_scale,
                 flow_permutation, flow_coupling, LU_decomposed,
                 flow_embed_dim, level):
        super().__init__()
        self.flow_coupling = flow_coupling

        self.actnorm = ActNorm2d(in_channels, actnorm_scale)

        # 2. permute
        if flow_permutation == "invconv":
            self.invconv = InvertibleConv1x1(in_channels,
                                             LU_decomposed=LU_decomposed)
            self.flow_permutation = \
                lambda z, logdet, rev: self.invconv(z, logdet, rev)
        elif flow_permutation == "shuffle":
            self.shuffle = Permute2d(in_channels, shuffle=True)
            self.flow_permutation = \
                lambda z, logdet, rev: (self.shuffle(z, rev), logdet)
        else:
            self.reverse = Permute2d(in_channels, shuffle=False)
            self.flow_permutation = \
                lambda z, logdet, rev: (self.reverse(z, rev), logdet)