### PyTorch 中 SSIM 实现 在 PyTorch 中，`torchvision.transforms.functional.ssim` 可用于计算两张图片之间的结构相似性指数 (SSIM)[^1]。然而，为了更灵活的应用场景以及将其作为损失函数的一部分，通常会自定义一个 `nn.Module` 类来封装 SSIM 计算逻辑。 下面是一个简单的 SSIM 损失类实现： ```python import torch import torch.nn as nn import torch.nn.functional as F class SSIMLoss(nn.Module): def __init__(self, window_size=11, size_average=True): super(SSIMLoss, self).__init__() self.window_size = window_size self.size_average = size_average self.channel = 1 self.window = create_window(window_size, self.channel) def forward(self, img1, img2): (_, channel, _, _) = img1.size() if channel == self.channel and self.window.data.type() == img1.data.type(): window = self.window else: window = create_window(self.window_size, channel).to(img1.device).type_as(img1) self.window = window self.channel = channel return 1 - ssim(img1, img2, window, self.window_size, channel, self.size_average) def gaussian(window_size, sigma): gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)]) return gauss/gauss.sum() def create_window(window_size, channel): _1D_window = gaussian(window_size, 1.5).unsqueeze(1) _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) window = _2D_window.expand(channel, 1, window_size, window_size).contiguous() return window def ssim(img1, img2, window, window_size, channel, size_average=True): mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) mu1_sq = mu1.pow(2) mu2_sq = mu2.pow(2) mu1_mu2 = mu1 * mu2 sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 C1 = 0.01 ** 2 C2 = 0.03 ** 2 ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) if size_average: return ssim_map.mean() else: return ssim_map.mean(1).mean(1).mean(1) ``` 此代码片段展示了如何创建一个可微分的 SSIM 损失模块，并通过卷积操作高效地计算局部均值和方差[^3]。注意这里采用了高斯加权窗口来进行平滑处理，这有助于捕捉图像的空间结构特征。 对于实际应用而言，可以直接调用上述定义好的 `SSIMLoss()` 来构建模型训练过程中所需的损失项。此外，如果只是单纯想要获取两张图间的 SSIM 值而不涉及反向传播，则只需关注正向传递部分即可。