Glow: Generative Flow with Invertible 1×1 Convolutions论文解析及实现(二)

Glow: Generative Flow with Invertible 1×1 Convolutions
代码github: https://github.com/rosinality/glow-pytorch添加链接描述

1 模型架构如下

1.1 左边图flow模型

Flow model
① ActNorm
② InvConv2dLU
③ AffineCoupling

1.2 右边模型结构Glow模型

Glow Model
Block (L-1)个
Flow K

2 Flow层

2.1 ActNorm

1）ActNorm 就是一个线性函数：

		# y=a*x +b 
        self.loc = nn.Parameter(torch.zeros(1, in_channel, 1, 1))   # b
        self.scale = nn.Parameter(torch.ones(1, in_channel, 1, 1))  #  a 

2） 参数初始化：
取第一批数据，计算均值方差，初始化 a 和b

 flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
            mean = ( flatten.mean(1) .unsqueeze(1) .unsqueeze(2).unsqueeze(3)
                .permute(1, 0, 2, 3)  )
            std = ( flatten.std(1)  .unsqueeze(1) .unsqueeze(2) .unsqueeze(3)
                .permute(1, 0, 2, 3))
            self.loc.data.copy_(-mean)
            self.scale.data.copy_(1 / (std + 1e-6))

3）log deteminant

log_abs = logabs(self.scale)
  logdet = height * width * torch.sum(log_abs)

4. 反函数

 output / self.scale - self.loc

2.2 Invertible 1*1 convolution

2.2.1使用LU 分解计算行列式

1）权重计算

        weight = (
            self.w_p
            @ (self.w_l * self.l_mask + self.l_eye)
            @ ((self.w_u * self.u_mask) + torch.diag(self.s_sign * torch.exp(self.w_s)))
        )

2）前向计算

  out = F.conv2d(input, weight)

3） 行列式计算

logdet = height * width * torch.sum(self.w_s)

4）反函数计算

 F.conv2d(output, weight.squeeze().inverse().unsqueeze(2).unsqueeze(3))

2.2.2 行列式直接计算方法

2.3 AffineCoupling

1）网络结构

        self.net = nn.Sequential(
            nn.Conv2d(in_channel // 2, filter_size, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(filter_size, filter_size, 1),
            nn.ReLU(inplace=True),
            ZeroConv2d(filter_size, in_channel if self.affine else in_channel // 2),
        )

2）前向公式

in_a, in_b = input.chunk(2, 1)
log_s, t = self.net(in_a).chunk(2, 1)
# s = torch.exp(log_s)
s = F.sigmoid(log_s + 2)
# out_a = s * in_a + t
out_b = (in_b + t) * s

torch.cat([in_a, out_b], 1)

3）log 行列式

 logdet = torch.sum(torch.log(s).view(input.shape[0], -1), 1)

3）反函数

out_a, out_b = output.chunk(2, 1)


log_s, t = self.net(out_a).chunk(2, 1)
# s = torch.exp(log_s)
s = F.sigmoid(log_s + 2)
# in_a = (out_a - t) / s
in_b = out_b / s - t

torch.cat([out_a, in_b], 1)

2.4 Flow 模型

1)前向计算及log行列式

out, logdet = self.actnorm(input)
out, det1 = self.invconv(out)
out, det2 = self.coupling(out)
logdet = logdet + det1+det2

2）反函数

        input = self.coupling.reverse(output)
        input = self.invconv.reverse(input)
        input = self.actnorm.reverse(input)

3 Glow 层

Flow * k
Block * (L-s)

3.1 Block

3.1.1 squeeze

        squeezed = input.view(b_size, n_channel, height // 2, 2, width // 2, 2)
        squeezed = squeezed.permute(0, 1, 3, 5, 2, 4)
        out = squeezed.contiguous().view(b_size, n_channel * 4, height // 2, width // 2)

3.1.2 step of flow

        for i in range(n_flow):
            self.flows.append(Flow(squeeze_dim, affine=affine, conv_lu=conv_lu))

3.1.3 split

            out, z_new = out.chunk(2, 1)
            mean, log_sd = self.prior(out).chunk(2, 1)
            log_p = gaussian_log_p(z_new, mean, log_sd)
            log_p = log_p.view(b_size, -1).sum(1)

3.1.4block 层

1）反函数

   input = output

        if reconstruct:
            if self.split:
                input = torch.cat([output, eps], 1)

            else:
                input = eps

        else:
            if self.split:
                mean, log_sd = self.prior(input).chunk(2, 1)
                z = gaussian_sample(eps, mean, log_sd)
                input = torch.cat([output, z], 1)

            else:
                zero = torch.zeros_like(input)
                # zero = F.pad(zero, [1, 1, 1, 1], value=1)
                mean, log_sd = self.prior(zero).chunk(2, 1)
                z = gaussian_sample(eps, mean, log_sd)
                input = z

        for flow in self.flows[::-1]:
            input = flow.reverse(input)

        b_size, n_channel, height, width = input.shape

        unsqueezed = input.view(b_size, n_channel // 4, 2, 2, height, width)
        unsqueezed = unsqueezed.permute(0, 1, 4, 2, 5, 3)
        unsqueezed = unsqueezed.contiguous().view(
            b_size, n_channel // 4, height * 2, width * 2
        )

        return unsqueezed

3.2 Glow层

1）前向计算

     self.blocks = nn.ModuleList()
        n_channel = in_channel
        for i in range(n_block - 1):
            self.blocks.append(Block(n_channel, n_flow, affine=affine, conv_lu=conv_lu))
            n_channel *= 2
        self.blocks.append(Block(n_channel, n_flow, split=False, affine=affine))

2）反函数

        for i, block in enumerate(self.blocks[::-1]):
            if i == 0:
                input = block.reverse(z_list[-1], z_list[-1], reconstruct=reconstruct)

            else:
                input = block.reverse(input, z_list[-(i + 1)], reconstruct=reconstruct)

        return input