检测异常图像中的异常是一项重要任务，尤其是在实时计算机视觉应用中。由于异常样本少，出现异常部位有一定随机性，因此目前在工业检测的普遍方法是使用可用的正常样本构建模型，以检测各种异常图像，而无需依赖真实的异常样本。这篇文章将构建一个简单的图像生成模型（Latent Diffusion Model，下文均称为LDM）进行尝试进行异常数据与正常数据的二分类判别，一起来看看吧！
MNIST数据集简介
MNIST数据集（Modified National Institute of Standards and Technology database）是一个广泛使用的手写数字数据库，主要用于训练各种图像处理系统。它是机器学习和深度学习的入门经典数据集之一。MNIST数据集图像数据是28x28的二维数组，每个元素是0到255之间的像素值，表示不同的灰度级别，每张图像表示一个手写数字（0到9）。总共包含70,000张图像，其中60,000张用于训练，10,000张用于测试。
它不是手写数字体分类数据集，为什么它能用来进行异常检测呢？
如果把一个数字类别设为正常类别，将其它类别设为异常数据，这样是不是就能训练模型检测这些异常数据了呢。
Latent Diffusion Model在异常检测的应用
如果大家对Latent Diffusion Model不太熟悉，可以到我的上篇博客Latent Diffusion Model是怎么完成图像快速生成的？潜在扩散模型（Latent Diffusion M - 掘金 (juejin.cn)上看下它的原理。来看下怎么把Latent Diffusion Model应用在异常检测上，先来看张流程图吧
图1 LDM在异常检测上的流程图
step1:训练LDM模型，输入正常类别训练数据，通过迭代一定次数的autoencoder及diffusion model训练后实现
step2:获得生成图像，输入测试数据，通过已训练好的LDM模型将测试图像生成为目标图像
step3:计算损失分数，通过计算原始图像与生成图像之间的差异实现损失分数的计算
看完这个流程图是不是感觉一下子明白了应用过程了呢？由此再延伸一下具体的异常检测实际应用，对于异常检测主要分为无监督检测、半监督检测、有监督检测，对于大多数情况下，异常样本收集非常困难，所以运用最广泛的方法为无监督检测。目前无监督检测在异常检测中的主流方法为：1）师生网络；2）生成式网络；3）Memory Bank；4）基于特征嵌入的方法。而本文中的方法属于生成式网络的方法，主要流程是通过重构测试样本到正常图像的分布，再通过计算目标图像与测试图像的差异进行计算损失，而有的生成模型甚至直接以是否能把测试样本重构为正常图像作为判别是否存在缺陷的标准，更为简单粗暴，但对于背景复杂的样本却不能起到很好地效果。对于师生网络，大家也可以看看我的上篇博客异常检测--基于特征提取的教师学生网络进行异常匹配原理篇（上）教师学生网络模型通常用于进行知识蒸馏任务，该任务的目的通常 - 掘金 (juejin.cn)进行学习。如果大家还想学习更多的异常检测主流方法，欢迎留言催更，你们的兴趣就是我的动力🧐🧐🧐。
Latent Diffusion Model搭建
数据集加载
数据预处理
transform = transforms.Compose([ transforms.ToTensor() ])
这里主要是把图像数据类型转换为神经网络处理所需的Tensor类型，如果到真实数据集可以加上图像增强的技术，如翻转，裁剪，灰度变换等方法
数据集加载
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
这里通过torchvision中的datasets函数直接下载数据集到data文件夹中，如果你已经有该数据可以修改root为文件夹路径，download设为False
读取指定类的数据
def get_indices_of_digit(dataset, digit):
indices = [i for i, (_, label) in enumerate(dataset) if label == digit]
return indices
train_indices = get_indices_of_digit(train_dataset, 0) # 仅包含数字0的训练集
train_dataset = Subset(train_dataset, train_indices)
这里主要完成的是对特定类别的数据进行读取，因为该方法只使用正类来进行检测异常数据，因此只需要读取一个正类即可
数据集按特定方式读入
train_loader = DataLoader(train_dataset, batch_size=512, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=512, shuffle=False)
将图像按512张图像每批的方式输入神经网络
网络搭建
class Autoencoder(nn.Module):
def __init__(self, latent_dim):
super(Autoencoder, self).__init__()
# Encoder
self.encoder = nn.Sequential(
nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.Flatten(),
nn.Linear(128*7*7, latent_dim)
)
# Decoder
self.decoder = nn.Sequential(
nn.Linear(latent_dim, 128*7*7),
nn.ReLU(),
nn.Unflatten(1, (128, 7, 7)),
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.ConvTranspose2d(64, 1, kernel_size=4, stride=2, padding=1),
nn.Sigmoid()
)
def encode(self, x):
return self.encoder(x)
def decode(self, z):
return self.decoder(z)
def forward(self, x):
z = self.encode(x)
x_reconstructed = self.decode(z)
return x_reconstructed
class LatentDiffusionModel(nn.Module):
def __init__(self, latent_dim, diffusion_steps=1000, beta_start=0.0001, beta_increment=0.0000199):
super(LatentDiffusionModel, self).__init__()
self.latent_dim = latent_dim
self.diffusion_steps = diffusion_steps
# 初始化 beta 序列
self.beta = torch.linspace(beta_start, beta_start + diffusion_steps * beta_increment, diffusion_steps)
# 计算 alpha 和 alpha_hat 序列
self.alpha = 1 - self.beta
self.alpha_hat = torch.cumprod(self.alpha, dim=0)
# 定义去噪网络
self.denoise_model = nn.Sequential(
nn.Linear(latent_dim, 512),
nn.ReLU(),
nn.Linear(512, 512),
nn.ReLU(),
nn.Linear(512, latent_dim)
)
def forward_diffusion(self, z, t):
noise = torch.randn_like(z)
hat_alpha_t = self.alpha_hat[t]
z_noisy = z * hat_alpha_t.sqrt() + noise * (1 - hat_alpha_t).sqrt()
return z_noisy
def reverse_diffusion(self, z, t):
beta_t = self.beta[t]
alpha_t = self.alpha[t]
hat_alpha_t = self.alpha_hat[t]
# 计算去噪网络输出
eps_theta = self.denoise_model(z)
# 反向扩散步骤计算公式
mean = (1 / alpha_t.sqrt()) * (z - (beta_t / (1 - hat_alpha_t).sqrt()) * eps_theta)
if t > 0:
z = mean + beta_t.sqrt() * torch.randn_like(z)
else:
z = mean # 最后一步不再添加噪声
return z
def sample(self, z_init, timesteps):
z = z_init
for t in reversed(range(timesteps)):
z = self.reverse_diffusion(z, t)
return z
构建一个autoencoder及diffusion网络，自动编码器通常由两部分组成：编码器（Encoder）和解码器（Decoder）。编码器负责将输入数据压缩成一个低维的潜在表示，而解码器则将这个潜在表示恢复成原始数据的近似版本。代码中使用卷积层、池化层、全连接层将数据压缩为latent_dim维度，再运用全连接层、池化层、卷积层将数据恢复成原始数据的近似版本，同时运用sigmod激活函数应用于该网络，因为该网络主要任务便是二分类问题，像素点的重构便相应的为二分类问题；diffusion网络与上篇原理篇的文章中的加噪与去噪步骤相同，但代码中仅使用了连接层进行重构图像，因为该数据集的简易性以及任务的不同，U-Net主要针对不同方面的输入，如文本、语义分割图、边缘图等。
损失函数设置
def reconstruction_loss(recon_x, x):
loss = F.mse_loss(recon_x, x, reduction='sum')
return loss
def kl_divergence(mu, logvar):
# KL 散度公式: D_KL[q(z|x) || p(z)] = -0.5 * sum(1 + logvar - mu^2 - exp(logvar))
kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return kl_loss
def vq_regularization(z_e, z_q):
# 计算两个损失
commitment_loss = F.mse_loss(z_e.detach(), z_q) # 使编码器输出接近量化向量
codebook_loss = F.mse_loss(z_e, z_q.detach()) # 使量化向量接近编码器输出
# 这里的损失通常以一定的权重相加
vq_loss = commitment_loss + codebook_loss
return vq_loss
分别定义了MSE、KL、VQ损失
- 均方误差（Mean Squared Error, MSE） ： 均方误差是最常用的损失函数之一，特别是在回归问题中。它计算预测值与真实值之间差的平方的平均值。Kullback-Leibler散度（Kullback-Leibler Divergence, KL Divergence） ： KL散度是衡量两个概率分布差异的一种方法，常用于衡量连续分布之间的差异。在机器学习中，KL散度可以用来衡量模型输出的概率分布与真实分布之间的差异。
- 向量量化（Vector Quantization, VQ）损失： 向量量化是一种将连续的输入数据映射到离散的表示空间的方法。VQ损失通常用于量化模型的性能，特别是在使用向量量化作为编码器的自动编码器中。VQ损失的计算方式是找到最接近输入数据的编码向量，并计算它们之间的差异。
MSE适用于需要预测连续值的问题，KL散度适用于需要比较概率分布差异的问题，而VQ损失适用于需要离散化表示的问题。
训练autoencoder
latent_dim = 128
num_epochs = 150
autoencoder = Autoencoder(latent_dim).cuda()
diffusion_model = LatentDiffusionModel(latent_dim).cuda()
ae_optimizer = Adam(autoencoder.parameters(), lr=1e-4)
diffusion_optimizer = Adam(diffusion_model.parameters(), lr=0.005)
torch.cuda.empty_cache()
# Train Autoencoder
for epoch in range(num_epochs):
for i, (images, _) in enumerate(train_loader):
images = images.cuda()
recon_images = autoencoder(images)
loss = vq_regularization(recon_images, images)
ae_optimizer.zero_grad()
loss.backward()
ae_optimizer.step()
print(f'loss:{loss}')
使用Adam优化器进行损失优化，同时使用VQ损失作为损失函数
之后的代码就是训练diffusion模型以及图像测试及保存，都是一些简单的代码，就不一一陈列了，如有需要可以私聊我获取。
测试结果
通过上面这些图可以发现，无论原始输入即原图是什么，生成图都会将其生成0，原图和生成图做差后得到的图片因此也会产生不同的差异。而原图是0的图像与生成图作差后得到的图像可以看到两图像差异小。
输出作差后的图像差异图，可以得到具体差异值，通过设置这个差异值为阈值不就可以完成对正常图像与异常图像的二分类问题了吗
总结
图像生成真的挺神奇的，通过图像重构的方式可以生成逼真、符合现实的图像，对此，它能否应用在更广泛的方面呢？比如，VR的应用，影视作品的生成，虚拟背景的嵌入，游戏开发等。再回味一下，新世纪的大门是不是已经被你敲开了呢，它不仅仅能推动工业的发展，还能够在教育、设计、医疗等多个领域发挥重要作用。接下来图像生成将发展如何，拭目以待吧。
