一、U-NET
- 论文题目:U-Net: Convolutional Networks for Biomedical Image Segmentation
- UNet 的体系结构基于编码器-解码器范式,其中编码器从输入图像中提取特征,解码器基于这些特征生成分割图。但是,UNet还集成了编码器和解码器之间的跳跃连接,以保留空间信息并提高分割精度。
- 先通过编码器进行特征提取,再通过解码器进行特征还原。
应用一Sepmark为例
Encoder
类是整个模型的核心,它负责将水印信息嵌入图像中。模型初始化时接受水印长度、模块中的卷积块数量、通道数以及注意力机制作为参数。主要操作:下采样——>上采样+信息嵌入+注意力块
1、下采样路径
self.conv1
: 初始卷积块,用于提取图像特征。self.down1
到self.down4
: 一系列下采样层,用于逐步减少图像的空间分辨率并增加通道数。
self.conv1 = ConvBlock(3, 16, blocks=blocks)
self.down1 = Down(16, 32, blocks=blocks)
self.down2 = Down(32, 64, blocks=blocks)
self.down3 = Down(64, 128, blocks=blocks)self.down4 = Down(128, 256, blocks=blocks)
其中下采样:
#Down 类:负责执行下采样操作,通过两次卷积块操作,第一次使用步长为2的卷积降采样,第二次卷积用于特征变换。
class Down(nn.Module):def __init__(self, in_channels, out_channels, blocks):super(Down, self).__init__()self.layer = torch.nn.Sequential(ConvBlock(in_channels, in_channels, stride=2),ConvBlock(in_channels, out_channels, blocks=blocks))def forward(self, x):return self. Layer(x)
2、上采样与水印嵌入路径
self.up3
到self.up0
: 上采样层,用于逐步恢复图像的空间分辨率。self.linear3
到self.linear0
: 线性层,用于将水印信息调整为适合嵌入的特征图尺寸。self.Conv_message3
到self.Conv_message0
: 卷积块,用于处理水印特征。self.att3
到 `self.att0: 注意力块,用于增强特征表示并融合水印信息。
self.up3 = UP(256, 128)
self.linear3 = nn.Linear(message_length, message_length * message_length)
self.Conv_message3 = ConvBlock(1, channels, blocks=blocks)
self.att3 = ResBlock(128 * 2 + channels, 128, blocks=blocks, attention=attention)self.up2 = UP(128, 64)
self.linear2 = nn.Linear(message_length, message_length * message_length)
self.Conv_message2 = ConvBlock(1, channels, blocks=blocks)
self.att2 = ResBlock(64 * 2 + channels, 64, blocks=blocks, attention=attention)self.up1 = UP(64, 32)
self.linear1 = nn.Linear(message_length, message_length * message_length)
self.Conv_message1 = ConvBlock(1, channels, blocks=blocks)
self.att1 = ResBlock(32 * 2 + channels, 32, blocks=blocks, attention=attention)self.up0 = UP(32, 16)
self.linear0 = nn.Linear(message_length, message_length * message_length)
self.Conv_message0 = ConvBlock(1, channels, blocks=blocks)
self.att0 = ResBlock(16 * 2 + channels, 16, blocks=blocks, attention=attention)
其中,上采样:
#UP 类:负责上采样操作,首先使用最近邻插值将输入特征图扩大两倍,然后通过一个卷积块处理放大的特征图。
class UP(nn.Module):def __init__(self, in_channels, out_channels):super(UP, self).__init__()self.conv = ConvBlock(in_channels, out_channels)def forward(self, x):x = F.interpolate(x, scale_factor=2, mode='nearest')return self.conv(x)
3、最终输出层
self.Conv_1x1
: 1x1 卷积层,用于生成最终的编码图像。
self.Conv_1x1 = nn.Conv2d(16 + 3, 3, kernel_size=1, stride=1, padding=0)
4、模型的前向传播方法
forward
方法定义了模型如何将水印信息嵌入图像中。它接收原始图像 x
和水印信息 watermark
作为输入,通过下采样路径提取特征,然后通过上采样路径逐步恢复图像分辨率,同时嵌入水印信息。
5、模型的输出
最终模型输出是嵌入水印后的图像,以及一个表示水印嵌入差距的 gap
。
from . import *#这段代码定义了一个名为 DW_Encoder 的深度学习模型类,它继承自 PyTorch 的 nn.Module 类。
# 这个模型主要设计用于图像隐写术,目的是在图像中嵌入水印(watermark)信息,同时尽量保持图像的视觉质量不被察觉。
# 模型结构采用了一种编码器-解码器架构,包含下采样(downsampling)和上采样(upsampling)路径,并且在编码过程中融入了注意力机制(通过可选的 attention 参数控制)。
class DW_Encoder(nn.Module):#初始化 (__init__):def __init__(self, message_length, blocks=2, channels=64, attention=None):# 接受多个参数,包括水印长度(message_length)、下采样块数(blocks)、通道数(channels)和注意力机制选项(attention)。super(DW_Encoder, self).__init__()# 定义了一系列卷积块(ConvBlock)、下采样层(Down)和上采样层(UP),以及用于融合水印信息的线性层(nn.Linear)和残差块(ResBlock),后者可以带有注意力机制。# 卷积块(ConvBlock):包含卷积层、激活函数(如ReLU)和可能的归一化层(如BatchNorm),用于特征提取和映射。self.conv1 = ConvBlock(3, 16, blocks=blocks)self.down1 = Down(16, 32, blocks=blocks)self.down2 = Down(32, 64, blocks=blocks)self.down3 = Down(64, 128, blocks=blocks)self.down4 = Down(128, 256, blocks=blocks)self.up3 = UP(256, 128)self.linear3 = nn.Linear(message_length, message_length * message_length)self.Conv_message3 = ConvBlock(1, channels, blocks=blocks)self.att3 = ResBlock(128 * 2 + channels, 128, blocks=blocks, attention=attention)self.up2 = UP(128, 64)self.linear2 = nn.Linear(message_length, message_length * message_length)self.Conv_message2 = ConvBlock(1, channels, blocks=blocks)self.att2 = ResBlock(64 * 2 + channels, 64, blocks=blocks, attention=attention)self.up1 = UP(64, 32)self.linear1 = nn.Linear(message_length, message_length * message_length)self.Conv_message1 = ConvBlock(1, channels, blocks=blocks)self.att1 = ResBlock(32 * 2 + channels, 32, blocks=blocks, attention=attention)self.up0 = UP(32, 16)self.linear0 = nn.Linear(message_length, message_length * message_length)self.Conv_message0 = ConvBlock(1, channels, blocks=blocks)self.att0 = ResBlock(16 * 2 + channels, 16, blocks=blocks, attention=attention)# 最后,定义一个1x1卷积层(Conv_1x1)来输出最终的编码图像,以及一些用于处理输入和水印数据的辅助函数。self.Conv_1x1 = nn.Conv2d(16 + 3, 3, kernel_size=1, stride=1, padding=0)self.message_length = message_length#在编码过程中融入了注意力机制(通过可选的 attention 参数控制)self.transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])def forward(self, x, watermark):#这是神经网络模型执行前向传播的核心逻辑。它接收两个输入参数:原始图像 x 和需要嵌入的水印信息 watermark。# 下采样路径(Encoder):首先,通过一系列称为Down的下采样模块(self.down1, self.down2, self.down3, self.down4)# 对输入图像x进行连续的下采样和特征提取,每次下采样都减少了空间分辨率但增加了通道数,生成了多个特征图d0, d1, d2, d3, d4。d0 = self.conv1(x)d1 = self.down1(d0)d2 = self.down2(d1)d3 = self.down3(d2)d4 = self.down4(d3)#上采样与信息嵌入路径(Decoder):然后,通过上采样模块 UP(self.up3, self.up2, self.up1, self.up0)# 以及线性层和卷积层将特征图逐步上采样回原始尺寸,同时在每个上采样步骤中,通过 linearX 和 Conv_messageX 将水印信息 watermark 扩展、调整尺寸并与当前级别的特征图拼接,# 然后通过注意力模块 attX 进一步融合信息。这里的操作旨在将水印以分散的方式嵌入到图像的不同尺度特征中。#这里,up3是一个上采样操作(由类UP定义),它接收来自下一层(第四层,d4)的特征图,# 并通过双线性插值(或其它指定的插值方式,本例中为nearest)将其尺寸放大一倍,之后可能还经过卷积层来提升特征维度,最终输出上采样的特征图u3。u3 = self.up3(d4)#这行代码将输入的水印信息通过一个全连接层(线性层)linear3进行变换,目的是调整水印的形状和维度,使其适合作为特征图与图像特征进行融合。expanded_message = self.linear3(watermark)#调整上一步输出的形状,使其成为一个四维张量,适合进一步处理作为图像特征。形状调整为批量大小(-1自动推断)、通道数1以及两个与水印长度相关的维度。expanded_message = expanded_message.view(-1, 1, self.message_length, self.message_length)#使用最近邻插值方法调整上一步得到的水印特征图尺寸,使之与第三层特征图d3的高和宽匹配,以便于后续的拼接操作。expanded_message = F.interpolate(expanded_message, size=(d3.shape[2], d3.shape[3]),mode='nearest')#通过卷积层Conv_message3对调整尺寸后的水印特征图进行卷积操作,这一步可以理解为对水印信息进行进一步的特征转换,使其更适应与图像特征的融合。expanded_message = self.Conv_message3(expanded_message)#这行代码将上采样的特征图u3、原始的第三层特征图d3以及经过处理后的水印特征图expanded_message沿着通道维度(dim=1)拼接起来,形成一个新的特征图,其中包含了原始图像信息和水印信息。u3 = torch.cat((d3, u3, expanded_message), dim=1)#最后,将拼接后的特征图输入到注意力模块att3中。# 注意力机制在这里的作用是自动学习并赋予不同区域或特征不同的权重,以突出重要信息并抑制无关信息,特别是在处理包含水印信息的特征时,有助于更好地融合和隐藏水印,同时保持图像质量。u3 = self.att3(u3)u2 = self.up2(u3)expanded_message = self.linear2(watermark)expanded_message = expanded_message.view(-1, 1, self.message_length, self.message_length)expanded_message = F.interpolate(expanded_message, size=(d2.shape[2], d2.shape[3]),mode='nearest')expanded_message = self.Conv_message2(expanded_message)u2 = torch.cat((d2, u2, expanded_message), dim=1)u2 = self.att2(u2)u1 = self.up1(u2)expanded_message = self.linear1(watermark)expanded_message = expanded_message.view(-1, 1, self.message_length, self.message_length)expanded_message = F.interpolate(expanded_message, size=(d1.shape[2], d1.shape[3]),mode='nearest')expanded_message = self.Conv_message1(expanded_message)u1 = torch.cat((d1, u1, expanded_message), dim=1)u1 = self.att1(u1)u0 = self.up0(u1)expanded_message = self.linear0(watermark)expanded_message = expanded_message.view(-1, 1, self.message_length, self.message_length)expanded_message = F.interpolate(expanded_message, size=(d0.shape[2], d0.shape[3]),mode='nearest')expanded_message = self.Conv_message0(expanded_message)u0 = torch.cat((d0, u0, expanded_message), dim=1)u0 = self.att0(u0)image = self.Conv_1x1(torch.cat((x, u0), dim=1))forward_image = image.clone().detach()'''read_image = torch.zeros_like(forward_image)for index in range(forward_image.shape[0]):single_image = ((forward_image[index].clamp(-1, 1).permute(1, 2, 0) + 1) / 2 * 255).add(0.5).clamp(0, 255).to('cpu', torch.uint8).numpy()im = Image.fromarray(single_image)read = np.array(im, dtype=np.uint8)read_image[index] = self.transform(read).unsqueeze(0).to(image.device)gap = read_image - forward_image'''gap = forward_image.clamp(-1, 1) - forward_imagereturn image + gap#Down 类:负责执行下采样操作,通过两次卷积块操作,第一次使用步长为2的卷积降采样,第二次卷积用于特征变换。
class Down(nn.Module):def __init__(self, in_channels, out_channels, blocks):super(Down, self).__init__()self.layer = torch.nn.Sequential(ConvBlock(in_channels, in_channels, stride=2),ConvBlock(in_channels, out_channels, blocks=blocks))def forward(self, x):return self.layer(x)#UP 类:负责上采样操作,首先使用最近邻插值将输入特征图扩大两倍,然后通过一个卷积块处理放大的特征图。
class UP(nn.Module):def __init__(self, in_channels, out_channels):super(UP, self).__init__()self.conv = ConvBlock(in_channels, out_channels)def forward(self, x):x = F.interpolate(x, scale_factor=2, mode='nearest')return self.conv(x)
在深度学习模型中,编号倒序(即从高到低的编号)通常用于表示网络层的连接顺序。在 Encoder
类中,编号倒序的原因是为了表示数据流经网络的顺序。每个编号对应于网络中的一个特定层或操作,编号越小,表示操作越先执行。以下是代码中编号倒序的解释:
-
下采样路径(Downsampling):
self.down1
,self.down2
,self.down3
,self.down4
:这些层用于逐步减少图像的空间分辨率并增加通道数。编号倒序表示数据流经这些层的顺序。 -
上采样路径(Upsampling):
self.up3
,self.up2
,self.up1
,self.up0
:这些层用于逐步恢复图像的空间分辨率。编号倒序表示数据流经这些层的顺序。 -
注意力块(ResBlock):
self.att3
,self.att2
,self.att1
,self.att0
:注意力块用于增强特征表示并融合水印信息。编号倒序表示这些块在上采样路径中的位置。 -
线性层(Linear):
self.linear3
,self.linear2
,self.linear1
,self.linear0
:线性层用于将水印信息调整为适合嵌入的特征图尺寸。 -
卷积块(ConvBlock):
self.Conv_message3
,self.Conv_message2
,self.Conv_message1
,self.Conv_message0
:卷积块用于处理水印特征。 -
最终输出层(Conv_1x1):
self.Conv_1x1
:1x1卷积层用于生成最终的编码图像。 -
辅助函数(transform):
self.transform
:用于处理输入和水印数据。
通过这种编号倒序,模型的设计者可以清晰地表示数据流经网络的顺序,即数据如何通过不同的层进行处理。这种顺序对于理解模型的前向传播过程至关重要,因为它决定了数据如何在网络中流动以及如何通过每一层。
二、注意力机制
与注意力机制相关的神经网络模块。这些模块可以用于增强模型的特征表示能力,特别是在图像识别、分割和生成等任务中。下面是对代码中每个部分的详细解释:[SEAttention效果>CBAM]
1. SEAttention (Squeeze-and-Excitation Attention)
- 功能:通过全局平均池化压缩空间信息到通道维度,然后通过两个全连接层(使用1x1卷积)调整通道权重,最后利用sigmoid函数输出每个通道的权重因子,重新调整原特征图的通道响应。
- 应用:强调重要的通道特征,抑制不重要的通道特征。【把重要的通道赋予大的权重,然后将这些通道以及权重去线性组合。】
- 即插即用的涨点模块之注意力机制(SEAttention)详解及代码,可应用于检测、分割、分类等各种算法领域-CSDN博客
#通过全局平均池化来压缩空间信息到通道维度,然后通过两个全连接层调整通道权重,最后利用sigmoid函数输出每个通道的权重因子,重新调整原特征图的通道响应。
class SEAttention(nn.Module):def __init__(self, in_channels, out_channels, reduction=8):super(SEAttention, self).__init__()self.se = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),nn.Conv2d(in_channels=in_channels, out_channels=out_channels // reduction, kernel_size=1, bias=False),nn.ReLU(inplace=True),nn.Conv2d(in_channels=out_channels // reduction, out_channels=out_channels, kernel_size=1, bias=False),nn.Sigmoid())def forward(self, x):x = self.se(x) * xreturn x
SE通道注意力机制模块-CSDN博客
2. ChannelAttention
- 功能:结合平均池化和最大池化后的通道注意力机制,通过这两个池化操作获取全局信息,然后通过两个卷积层调整通道权重,最后利用sigmoid激活函数输出注意力图。
- 应用:同时考虑平均和最大池化的信息,以更全面地捕捉通道间的关系。
#结合了平均池化和最大池化后的通道注意力机制,通过这两个池化操作获取全局信息,然后通过两个卷积层调整通道权重,最后利用sigmoid激活函数输出注意力图。
class ChannelAttention(nn.Module):def __init__(self, in_channels, out_channels, reduction=8):super(ChannelAttention, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))self.max_pool = nn.AdaptiveMaxPool2d((1, 1))self.fc = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels // reduction, kernel_size=1, bias=False),nn.ReLU(inplace=True),nn.Conv2d(in_channels=out_channels // reduction, out_channels=out_channels, kernel_size=1, bias=False))self.sigmoid = nn.Sigmoid()def forward(self, x):avg_out = self.fc(self.avg_pool(x))max_out = self.fc(self.max_pool(x))out = avg_out + max_outreturn self. Sigmoid(out)
3. SpatialAttention
- 功能:利用平均池化和最大池化来分别捕捉空间上下文信息,然后通过一个卷积层融合这两种信息并使用sigmoid激活函数输出空间注意力图,强调重要的空间位置。
- 应用:强调图像中的重要区域,抑制不重要的区域。
#利用平均池化和最大池化来分别捕捉空间上下文信息,然后通过一个卷积层融合这两种信息并使用sigmoid激活函数输出空间注意力图,强调重要的空间位置。
class SpatialAttention(nn.Module):def __init__(self, kernel_size=7):super(SpatialAttention, self).__init__()self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)self.sigmoid = nn.Sigmoid()def forward(self, x):avg_out = torch.mean(x, dim=1, keepdim=True)max_out, _ = torch.max(x, dim=1, keepdim=True)x = torch.cat([avg_out, max_out], dim=1)x = self.conv1(x)return self.sigmoid(x)
4. CBAMAttention
- 功能:结合上述的ChannelAttention和SpatialAttention,先计算通道注意力,再计算空间注意力,最后将两者相乘以指导特征图的注意力分配。
- 应用:同时考虑通道和空间的注意力,以更精细地控制特征图的加权。
#结合了上述的ChannelAttention和SpatialAttention,先计算通道注意力,再计算空间注意力,最后将两者相乘以指导特征图的注意力分配。
class CBAMAttention(nn.Module):def __init__(self, in_channels, out_channels, reduction=8):super(CBAMAttention, self).__init__()self.ca = ChannelAttention(in_channels=in_channels, out_channels=out_channels, reduction=reduction)self.sa = SpatialAttention()def forward(self, x):x = self.ca(x) * xx = self.sa(x) * xreturn x
5、ECAAttention
ECANet的核心思想是在卷积操作中引入通道注意力机制,以提升特征表示的能力。通道注意力机制旨在自适应地调整通道特征的权重,从而使网络能够更好地关注重要特征并抑制不重要特征。通过这种机制,ECANet有效地增强了网络的表征能力,同时避免了增加过多的参数和计算成本。
嵌入式通道注意力模块是ECANet的扩展部分,将通道注意力机制嵌入到卷积层中,从而在卷积操作中引入通道关系。具体而言,在卷积操作中,将输入特征图划分为多个子特征图,然后在每个子特征图上进行卷积操作,并在操作过程中引入通道注意力。最后,将这些子特征图合并,得到最终的输出特征图。这种设计有效地减少了计算成本,并保持了网络的高效性。
import torch
import torch.nn as nn
import mathclass ECA(nn.Module):def __init__(self, in_channel, gamma=2, b=1):super(ECA, self).__init__()k = int(abs((math.log(in_channel,2)+b)/gamma))kernel_size = k if k % 2 else k+1padding = kernel_size//2self.pool = nn.AdaptiveAvgPool2d(output_size=1)self.conv = nn.Sequential(nn.Conv1d(in_channels=1, out_channels=1, kernel_size=kernel_size, padding=padding, bias=False),nn.Sigmoid())def forward(self,x):out=self.pool(x)out=out.view(x.size(0), 1, x.size(1))out=self.conv(out)out=out.view(x.size(0), x.size(1), 1, 1)return out*x