Flow

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

# 1. Coupling Layer: 采用 RealNVP 风格的 Affine Coupling
class AffineCoupling(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.scale_net = nn.Sequential(
            nn.Conv2d(in_channels // 2, in_channels, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels, in_channels // 2, kernel_size=3, padding=1),
            nn.Tanh()  # 输出 scale
        )
        self.translate_net = nn.Sequential(
            nn.Conv2d(in_channels // 2, in_channels, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels, in_channels // 2, kernel_size=3, padding=1)
        )
    
    def forward(self, x, reverse=False):
        x1, x2 = x.chunk(2, dim=1)  # Split channels
        s = self.scale_net(x1)
        t = self.translate_net(x1)
        if not reverse:
            x2 = x2 * torch.exp(s) + t
        else:
            x2 = (x2 - t) * torch.exp(-s)
        return torch.cat([x1, x2], dim=1), torch.sum(s, dim=[1, 2, 3])  # 返回 log-determinant

# 2. Normalizing Flow Model
class RealNVP(nn.Module):
    def __init__(self, in_channels, num_flows=4):
        super().__init__()
        self.flows = nn.ModuleList([AffineCoupling(in_channels) for _ in range(num_flows)])

    def forward(self, x):
        log_det_jacobian = 0
        for flow in self.flows:
            x, log_det = flow(x, reverse=False)
            log_det_jacobian += log_det
        return x, log_det_jacobian

    def reverse(self, z):
        for flow in reversed(self.flows):
            z, _ = flow(z, reverse=True)
        return z

# 3. 训练函数
def train(model, data_loader, optimizer, num_epochs=10):
    model.train()
    for epoch in range(num_epochs):
        for x in data_loader:
            x = x.to(device)
            z, log_det = model(x)
            log_p_z = -0.5 * torch.sum(z ** 2, dim=[1, 2, 3])  # 假设标准正态分布
            loss = -torch.mean(log_p_z + log_det)  # 最大化对数似然
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch+1}, Loss: {loss.item()}")

# 4. 生成图像
def generate_samples(model, num_samples=16):
    z = torch.randn(num_samples, 3, 32, 32).to(device)  # 采样标准正态分布
    samples = model.reverse(z)  # 通过逆变换生成
    return samples

# 初始化模型与优化器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RealNVP(in_channels=3, num_flows=6).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 训练示例
# data_loader = ...  # 假设我们有一个图像数据集的 DataLoader
# train(model, data_loader, optimizer, num_epochs=10)

# 生成图像示例
# samples = generate_samples(model)
# visualize(samples)  # 假设有一个可视化函数