cm-diff：解析红外与可见光图像双向跨模态翻译扩散模型的单一生成网络实现

1次阅读

没有评论

共计 1527 个字符，预计需要花费 4 分钟才能阅读完成。

跨模态图像翻译任务在计算机视觉领域具有重要应用价值，特别是在红外与可见光图像的相互转换方面。传统方法通常采用多网络架构，即针对不同方向的翻译任务训练独立的网络模型。这种方案存在明显局限性：

资源消耗大 ：需要训练和维护多个网络，显存和计算资源需求成倍增加
参数冗余 ：不同网络可能学习到相似的底层特征表示
一致性差 ：双向翻译结果难以保持语义和风格的一致性

目前主流的生成模型主要包括 GAN、VAE 和扩散模型。在跨模态翻译任务中，它们各有优劣：

GAN：
优点：生成图像质量高，细节丰富
缺点：训练不稳定，容易出现模式崩溃
VAE：
优点：训练稳定，有理论保证
缺点：生成图像通常较模糊
扩散模型 ：
优点：训练稳定性好，生成质量高
缺点：推理速度较慢

cm-diff 选择扩散模型作为基础架构，通过创新设计克服了传统方法的局限性。

扩散模型通过逐步添加和去除噪声来实现图像生成。在跨模态翻译任务中，这一过程可以表示为：

前向过程：逐步向输入图像添加高斯噪声
反向过程：从噪声中重建目标模态图像

数学表示为：

q(x_t|x_{t-1}) = N(x_t; \sqrt{1-\beta_t}x_{t-1}, \beta_tI)

cm-diff 的核心创新在于使用单一网络处理双向翻译。关键技术包括：

模态条件输入 ：在输入中加入模态标签作为条件
共享特征提取 ：底层特征提取器在不同模态间共享
自适应归一化 ：根据模态条件调整归一化参数

损失函数设计 ：
结合 L1 损失和感知损失
加入对抗损失提高生成质量
噪声调度策略 ：
采用余弦调度控制噪声添加过程
在训练后期减少噪声强度

以下是核心模块的 PyTorch 实现：

import torch
import torch.nn as nn

class CM_Diff(nn.Module):
    def __init__(self, in_channels=3, out_channels=3, n_feat=64):
        super().__init__()
        # 共享的编码器部分
        self.encoder = nn.Sequential(nn.Conv2d(in_channels, n_feat, 3, padding=1),
            nn.GroupNorm(8, n_feat),
            nn.SiLU(),
            # 更多层...
        )

        # 模态条件处理
        self.modal_embed = nn.Embedding(2, n_feat)

        # 解码器部分
        self.decoder = nn.Sequential(
            # 解码层...
            nn.Conv2d(n_feat, out_channels, 3, padding=1)
        )

    def forward(self, x, modal):
        # modal: 0 表示红外 -> 可见光，1 表示可见光 -> 红外
        h = self.encoder(x)
        modal_emb = self.modal_embed(modal).unsqueeze(-1).unsqueeze(-1)
        h = h + modal_emb
        return self.decoder(h)

我们在标准数据集上进行了全面评估：