normalizing_flow.py

# Purpose: Invertible Normalizing Flow (RealNVP) to mathematically guarantee exact forward and inverse transformations of the latent space without data loss. Input: 3D latent tensor. Output: Transformed 3D latent tensor.
import torch
import torch.nn as nn

class ConvNet3D(nn.Module):
    def __init__(self, in_channels, out_channels, hidden_channels=256):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv3d(in_channels, hidden_channels, kernel_size=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv3d(hidden_channels, out_channels, kernel_size=1)
        )
        
        # Initialize final layer to zero so the flow starts as an identity function
        nn.init.zeros_(self.net[-1].weight)
        nn.init.zeros_(self.net[-1].bias)

    def forward(self, x):
        return self.net(x)


class AffineCouplingLayer3D(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.split_dim = 1 # Split along channel dimension
        assert channels % 2 == 0, "Channels must be divisible by 2 for coupling layer"
        half_channels = channels // 2
        
        # s (scale) and t (translation) networks
        self.s_net = ConvNet3D(half_channels, half_channels)
        self.t_net = ConvNet3D(half_channels, half_channels)

    def forward(self, x, reverse=False):
        x1, x2 = torch.chunk(x, 2, dim=self.split_dim)
        
        # Bound scale to [-1, 1] to prevent torch.exp() from exploding to infinity
        s = torch.tanh(self.s_net(x1))
        t = self.t_net(x1)

        if not reverse:
            # Forward: y1 = x1, y2 = x2 * exp(s) + t
            y1 = x1
            y2 = x2 * torch.exp(s) + t
        else:
            # Inverse: x1 = y1, x2 = (y2 - t) * exp(-s)
            y1 = x1
            y2 = (x2 - t) * torch.exp(-s)
            
        return torch.cat([y1, y2], dim=self.split_dim)


class CyLNormalizingFlow(nn.Module):
    def __init__(self, channels, num_layers=3):
        super().__init__()
        self.layers = nn.ModuleList([AffineCouplingLayer3D(channels) for _ in range(num_layers)])

    def forward(self, x, reverse=False):
        if not reverse:
            for layer in self.layers:
                x = layer(x, reverse=False)
                # Flip channels so different halves get transformed
                x = torch.flip(x, dims=[1])
        else:
            for layer in reversed(self.layers):
                x = torch.flip(x, dims=[1])
                x = layer(x, reverse=True)
        return x