图像分割UNet、生成模型SD及IP-Adapter

文章目录

• 摘要
• abstract
• 1.UNet网络
• • [1.1 UNet损失函数](#1.1 UNet损失函数)
  • [1.2 UNet结构](#1.2 UNet结构)
  • [1.3 UNet代码实现](#1.3 UNet代码实现)
• [2.Stable Diffusion](#2.Stable Diffusion)
• • [2.1 扩散模型基础（DM)](#2.1 扩散模型基础（DM))
  • [2.2 潜在扩散模型（LDM)](#2.2 潜在扩散模型（LDM))
  • [2.3 Encoder/Decoder](#2.3 Encoder/Decoder)
  • [2.4 代码示例](#2.4 代码示例)
• 3.IP-Adapter
• • [3.1 理论分析](#3.1 理论分析)
• 4.总结
• 参考文献

摘要

本周学习了UNet、Stable Diffusion及IP-Adapter三种计算机视觉模型的原理和实现。UNet是一种用于语义分割的神经网络，采用编码-解码结构，能够进行像素级分类，并通过带权交叉熵损失函数优化边界分割。Stable Diffusion是一种基于潜在扩散模型（LDM）的文本到图像生成方法，通过变分自编码器（VAE）降维并使用交叉注意力UNet进行去噪，使得图像生成更加高效。IP-Adapter针对Stable Diffusion的局限性，提出了解耦交叉注意力机制，通过额外的注意力层引入图像特征，以更精确地控制生成结果，提高图像一致性和细节保真度。

abstract

This week, we learned the principle and implementation of three computer vision models: UNet, Stable Diffusion and IP-Adapter. UNet is a neural network for semantic segmentation, which adopts a coding-decoding structure, is capable of pixel-level classification, and optimizes boundary segmentation by weighted cross-entropy loss function. Stable Diffusion is a text-to-image generation method based on potential diffusion model (LDM), which makes image generation more efficient by variational autoencoder (VAE) dimensionality reduction and denoising using cross-attention UNet. Aiming at the limitation of Stable Diffusion, IP-Adapter proposes an decouple cross-attention mechanism, which introduces image features through additional attention layers to control the generated results more precisely and improve image consistency and detail fidelity.

1.UNet网络

UNet是一个语义分割模型，主要执行过程与其他语义分割模型类似。与CNN不同之处是CNN是图像级的分类，而UNet是像素级分类，其输出的是每个像素点的类别。

1.1 UNet损失函数

主要组成部分：softmax激活函数+带权的交叉熵损失函数+权重计算函数
softmax激活函数

softmax激活函数将每个像素点的输入特征与权重做非线性叠加。每个像素点金经过softmax的处理后产生的输出值个数等于标签里的类别数。softmax将每个像素点的输出值变换成值为正和为1的概率分布，从而得到每个像素点中每个类的置信度。
S i = e V i ∑ i C e V i S_i=\frac{e^{V_i}}{\sum_i^Ce^{V_i}} Si=∑iCeVieVi
交叉熵损失函数

用来衡量两个概率分布差异的测量函数
L = − ∑ c = 1 M y c l o g ( p c ) L=-\sum_{c=1}^My_clog(p_c) L=−c=1∑Myclog(pc)

在上式中y~c~表示样本的真实分布，值要么为0，要么为1，P~c~表示样本的预测分布。

在UNet论文中使用带边界权重的交叉熵损失函数：
E = ∑ x ∈ Ω w ( x ) log ⁡ ( p ℓ ( x ) ( x ) ) E=\sum_{\mathbf{x}\in\Omega}w(\mathbf{x})\log(p_{\ell(\mathbf{x})}(\mathbf{x})) E=x∈Ω∑w(x)log(pℓ(x)(x))

其中pl(x)是点x在对应的lablel给出的那个类别输出的激活值。
权重计算函数
w ( x ) = w c ( x ) + w 0 ⋅ exp ⁡ ( − ( d 1 ( x ) + d 2 ( x ) ) 2 2 σ 2 ) w(\mathbf{x})=w_c(\mathbf{x})+w_0\cdot\exp\left(-\frac{(d_1(\mathbf{x})+d_2(\mathbf{x}))^2}{2\sigma^2}\right) w(x)=wc(x)+w0⋅exp(−2σ2(d1(x)+d2(x))2)

权重的作用：权重可以调整图像中某个区域的重要程度，在计算损失的过程中（医学细胞分割），给两个细胞重合的边缘部分增加了损失的权重，让网络更加关注这类重合的边缘信息。

1.2 UNet结构

U-Net 主要由 编码器（Contracting Path） 和 解码器（Expanding Path） 组成，结构如下：
输入： 大小为256256的图片
输出： 相同大小的分割掩码
编码器： 提取高级特征，同时降低空间分辨率（采用两次3 3卷积+ReLU激活函数，然后使用2*2最大池化进行下采样（特征图尺寸减半）），每次下采样时，通道数翻倍，用于学习更多特征。
解码器： 恢复空间分辨率，生成分割掩码（上采样[转置卷积/双线性插值]+跳跃连接[通道拼接-编码过程中丢失的高分辨率信息，同时恢复位置信息]）

1.3 UNet代码实现

class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super(UNet, self).__init__()

        # 编码部分
        self.conv1 = self.double_conv(in_channels, 64)
        self.conv2 = self.double_conv(64, 128)
        self.conv3 = self.double_conv(128, 256)
        self.conv4 = self.double_conv(256, 512)
        self.conv5 = self.double_conv(512, 1024)

        # 池化层
        self.pool = nn.MaxPool2d(2)

        # 解码部分
        self.up6 = self.up_conv(1024, 512)
        self.conv6 = self.double_conv(1024, 512)

        self.up7 = self.up_conv(512, 256)
        self.conv7 = self.double_conv(512, 256)

        self.up8 = self.up_conv(256, 128)
        self.conv8 = self.double_conv(256, 128)

        self.up9 = self.up_conv(128, 64)
        self.conv9 = self.double_conv(128, 64)

        # 最终 1x1 卷积
        self.final = nn.Conv2d(64, out_channels, kernel_size=1)

    def double_conv(self, in_channels, out_channels):
        """ 两个 3x3 卷积 """
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

    def up_conv(self, in_channels, out_channels):
        """ 上采样 + 转置卷积 """
        return nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)

    def forward(self, x):
        # 编码
        c1 = self.conv1(x)
        p1 = self.pool(c1)

        c2 = self.conv2(p1)
        p2 = self.pool(c2)

        c3 = self.conv3(p2)
        p3 = self.pool(c3)

        c4 = self.conv4(p3)
        p4 = self.pool(c4)

        c5 = self.conv5(p4)  # 最底层

        # 解码
        u6 = self.up6(c5)
        merge6 = torch.cat([u6, c4], dim=1)
        c6 = self.conv6(merge6)

        u7 = self.up7(c6)
        merge7 = torch.cat([u7, c3], dim=1)
        c7 = self.conv7(merge7)

        u8 = self.up8(c7)
        merge8 = torch.cat([u8, c2], dim=1)
        c8 = self.conv8(merge8)

        u9 = self.up9(c8)
        merge9 = torch.cat([u9, c1], dim=1)
        c9 = self.conv9(merge9)

        output = self.final(c9)
        return output

# 创建 U-Net 模型
model = UNet()
print(model)

2.Stable Diffusion

stable diffusion是一种潜在扩散模型的文本-图像生成模型。

2.1 扩散模型基础（DM)

回顾前两周学习的扩散模型（DM是一种概率生成模型，主要用于图像、音频、视频的生成任务。）

核心思想：
前向扩散： 将数据逐步加入高斯噪声，直至变为纯噪声
反向去噪： 训练一个神经网络（通常是U-Net）学习如何逐步去噪，从纯噪声中恢复清晰的图像。

2.2 潜在扩散模型（LDM)

传统扩散模型（如DDPM)在像素空间进行去噪，计算成本高。LDM提出了新的优化策略：
降维： 先将图像压缩到潜在空间，然后在该空间进行扩散和去噪。
计算量减少： 降低了计算需求，同时保持高质量生成。

LDM主要组成部分

1. VAE(变分自编码器)：负责将原始图像压缩到潜在空间，降低计算复杂度。生成完图片后，在通过VAE解码还原到像素空间。
2. U-Net(去噪神经网络)：采用交叉注意力（Cross-Attention），根据文本信息进行图像去噪。其结构类似于UNet(上采样+下采样)，用于处理多尺度信息。

交叉注意力的计算方式
A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V \mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=softmax(dk QKT)V

Q来自Image embedding,K和V来自condition embedding,通过Context的指导，来调节image中各个位置的重要性。

具体过程：

1. 文本提示通过CLIP文本编码器转换为嵌入向量
2. U-Net计算Q,文本嵌入提供K和V
3. 交叉注意力计算Q和K之间的匹配度，得到文本对图像特征的影响程度
4. V通过注意力权重调整Q,实现文本对图像去噪的控制
5. 交叉注意力输出的新特征被融合到U-Net进行进一步去噪
   

3.文本编辑器(T5或CLIP)：将文本转化为潜在表示，让模型能理解Prompt。SD默认使用CLIP Text Encoder,也可以支持T5。

2.3 Encoder/Decoder

其由多个resnet block组成

2.4 代码示例

下面是使用 交叉注意力U-Net 进行文本引导去噪的PyTorch代码框架

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

class CrossAttention(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.to_q = nn.Linear(dim, dim, bias=False)
        self.to_k = nn.Linear(dim, dim, bias=False)
        self.to_v = nn.Linear(dim, dim, bias=False)
        self.scale = dim ** -0.5

    def forward(self, image_features, text_features):
        Q = self.to_q(image_features)
        K = self.to_k(text_features)
        V = self.to_v(text_features)

        attn = F.softmax(torch.matmul(Q, K.transpose(-2, -1)) * self.scale, dim=-1)
        return torch.matmul(attn, V)

class UNetWithCrossAttention(nn.Module):
    def __init__(self, image_dim, text_dim):
        super().__init__()
        self.encoder = nn.Conv2d(image_dim, 64, kernel_size=3, padding=1)
        self.cross_attention = CrossAttention(64)
        self.decoder = nn.Conv2d(64, image_dim, kernel_size=3, padding=1)

    def forward(self, image, text_embedding):
        x = self.encoder(image)  # 提取图像特征
        x = self.cross_attention(x, text_embedding)  # 交叉注意力
        x = self.decoder(x)  # 解码
        return x

# 示例输入
image = torch.randn(1, 3, 64, 64)  # 随机噪声图像
text_embedding = torch.randn(1, 10, 64)  # 文本嵌入（来自 Transformer）

# 运行 U-Net 去噪
unet = UNetWithCrossAttention(image_dim=3, text_dim=10)
output = unet(image, text_embedding)
print(output.shape)

3.IP-Adapter

对于先前的文本到图像生成主要是从CLIP图像编码器中提取的图像特征通过可训练的网络映射到新的特征，然后与文本特征进行连接。通过替换原始文本特征，将合并后的特征输入到扩散模型的UNet中，指导图像生成。上述所面临的问题是生成的图像只是部分忠实于提示的图像。结果往往不如经过微调的图像提示模型。

针对上述问题分析《IP-Adapter: Text Compatible Image Prompt Adapter for Text-to-Image Diffusion Models》文中指出主要问题在于文本到图像扩散模型的交叉注意模块。交叉注意力层的k和v得值训练以适应文本特征。这使得图像特征和文本特征合并到交叉注意层只完成了图像特征与文本特征的对齐，上述方法可能会遗漏一些特定于图像得信息，最终导致仅与参考图像进行粗粒度得可控生成。文中采用文本特征和图像特征采用解耦得交叉注意机制。对于扩散模型UNet中的每一个交叉注意层，只针对图像特征增加一个额外的交叉注意层。训练过程中，只训练新的交叉注意层的参数，原始的UNet模型保持冻结。

上述是具有解耦交叉注意策略的IP-Adapter的整体架构。只有新添加的红色模块被训练，而预训练的文本到图像模型被冻结。

通过新增的Cross-Attention模块，参考图的特征会经过W~k~,W~v~权重，映射为K，V向量，噪声图像特征经过W~q~映射为查询向量Q.在训练时，只调整K，V，Q的权重即可。

3.1 理论分析

图像特征是把具有解耦的交叉注意力的适配模块集成到训练好的UNet中，而原始SD模型中，文本特征也是通过交叉注意力层将文本条件注入到UNet中，查询特征Z和文本特征C~s~，交叉注意力的输出Z'为：
Z ′ = A t t e n t i o n ( Q , K , V ) = S o f t m a x ( Q K T d ) V Z^{\prime}=\mathrm{Attention}(Q,K,V)=\mathrm{Softmax}(\frac{QK^T}{\sqrt{d}})V Z′=Attention(Q,K,V)=Softmax(d QKT)V

其中， Q = Z W q , K = c s W k , V = c s W v Q=ZW_{q},K=c_{s}W_{k},V=c_{s}W_{v} Q=ZWq,K=csWk,V=csWv

此时图像特征注入模型有两种方法，第一种是将图像特征和文本特征进行串联，然后输入到注意力交叉层，但是不够有效；另外一种，解耦的交叉注意力机制，将文本和图像特征的交叉注意力层分开，即在UNet的CA-Down/Up-2D模块中的每个Cross-Attention层，新增一个Cross-Attention层，这个新增的层用于将图像也特征插入。假设图像特征为c~f~,新增的Cross-Attention层的输出Z''为：
Z ′ ′ = A t t e n t i o n ( Q , K ′ , V ′ ) = S o f t m a x ( Q ⋅ ( K ′ ) T d ) V ′ Z^{\prime\prime}=\mathrm{Attention}(Q,K^{\prime},V^{\prime})=\mathrm{Softmax}(\frac{Q\cdot(K^{\prime})^T}{\sqrt{d}})V^{\prime} Z′′=Attention(Q,K′,V′)=Softmax(d Q⋅(K′)T)V′

其中， Q = Z W q , K ′ = c f W k ′ , V ′ = c f V v ′ Q=ZW_{q},K^{\prime}=c_{f}W_{k}^{\prime},V^{\prime}=c_{f}V_{v}^{\prime} Q=ZWq,K′=cfWk′,V′=cfVv′

K'和V'是来自参考图的k向量和v向量，查询向量Q与文本的查询矩阵是相同的，为了加快收敛，W~k~'和W~v~'是由W~k~,W~v~初始化。

最终结果将文本的Cross-Attention 输出Z'和图像的Cross-Attention输出的Z''相加 Z ′ ′ = S o f t m a x ( Q K T d ) V + λ ⋅ S o f t m a x ( Q ⋅ ( K ′ ) T d ) V ′ Z^{\prime\prime}=\mathrm{Softmax}(\frac{QK^T}{\sqrt{d}})V+\lambda\cdot\mathrm{Softmax}(\frac{Q\cdot(K^{\prime})^T}{\sqrt{d}})V^{\prime} Z′′=Softmax(d QKT)V+λ⋅Softmax(d Q⋅(K′)T)V′

由于冻结了最初UNet模型，在上述解耦交叉注意力中，只有W~k~'和W~v~'训练得到调整。

4.总结

UNet通过特征提取、下采样和上采样实现高精度语义分割，并广泛应用于医学图像处理。Stable Diffusion通过扩散模型在潜在空间进行去噪生成，提高了图像生成的效率和质量。IP-Adapter针对扩散模型的文本引导问题，提出解耦交叉注意力机制，使图像提示更加准确，有效提升了图像生成的一致性。

参考文献

[1] IP-Adapter: Text Compatible Image Prompt Adapter forText-to-Image Diffusion Models https://arxiv.org/pdf/2308.06721.pdf