文章目录
摘要
本周学习了潜在扩散模型的概念，潜在扩散模型的实现是很简单的，之后学习了变分自编码器(VAE)。同时解决了之前的RCNN目标检测的模型，但是效果很差。
Abstract
This week, I studied the concept of latent diffusion models. The implementation of latent diffusion models is quite simple. Afterwards, I learned about variational autoencoders (VAE). At the same time, I worked on the previous RCNN object detection model, but the results were very poor.
1 潜在扩散模型
传统扩散模型在像素空间直接操作的特性导致其训练和推理过程计算成本极高，限制了其在高分辨率图像生成中的广泛应用。为了解决这一问题，Latent Diffusion Model（潜在扩散模型，LDMs）应运而生，它通过在潜在空间中运行扩散过程，显著降低了计算复杂度，同时保留甚至提升了生成质量。
1.1 传统扩散模型
传统扩散模型的核心思想是通过一个逐步加噪的前向过程和一个逐步去噪的后向过程来学习数据的分布。具体来说，前向过程将原始数据（如图像）逐步添加高斯噪声，直至接近纯噪声分布；后向过程则通过训练一个神经网络（通常是 UNet）预测每一步的噪声，从而逆向重建数据。这种方法在理论上等价于最大似然估计，能够有效避免 GAN 的模式崩塌问题，并生成多样性更高的样本。
然而，传统 DMs 的一个显著缺点是其直接在高维像素空间中操作。以一张 256×256 的 RGB 图像为例，其维度高达 196,608（256×256×3），这意味着每次前向和后向步骤都需要处理海量数据。训练一个强大的像素空间 DM 通常需要数百个 GPU 天（例如，文献中提到 150-1000 个 V100 GPU 天），而推理过程也因需要数百到上千次顺序评估而变得昂贵。这种高计算成本不仅限制了模型的可扩展性，还对资源有限的研究者构成了门槛。
1.2 潜在扩散模型
LDM 的核心创新在于将扩散过程从像素空间转移到潜在空间（Latent Space），通过一个预训练的自动编码器（Autoencoder）将高维图像数据压缩为低维表示，再在此低维空间中进行扩散建模。以下是其工作原理的分解：
感知压缩（Perceptual Compression）
LDM 首先利用一个自动编码器将图像 ( x ∈ R H × W × 3 ) (x\in R^{H\times W \times 3}) (x∈RH×W×3) 编码为潜在表示 ( z = θ ( x ) ∈ R h × w × c ) (z= \theta(x)\in R^{h\times w \times c}) (z=θ(x)∈Rh×w×c)，其中 ( h = H / f ) 、 ( h = H / f ) ( ( f ) ) (h=H/f)、(h=H/f)((f)) (h=H/f)、(h=H/f)((f))为下采样因子，通常取 ( 2 m , m ∈ N ) (2^m,m\in N) (2m,m∈N)。解码器 ( D ) (\mathcal{D}) (D) 则将 x ^ = D ( z ) \hat{x}=\mathcal{D}(z) x^=D(z)重建为图像
自动编码器的训练目标是实现感知等价性，即 ( x ^ ) (\hat{x}) (x^)在感知上接近 ( x ) (x) (x)，而非像素级完全一致。为此，训练结合了感知损失（Perceptual Loss）和对抗损失（Adversarial Loss），确保重建图像保持局部真实感并避免模糊。
为了控制潜在空间的分布，LDM 引入两种正则化方法：KL 正则化（类似于 VAE，限制 ( z ) (z) (z) 接近标准正态分布）和 VQ 正则化（通过向量量化层离散化 ( z ) ( z ) (z)。这些正则化确保潜在空间的稳定性和可控性。
潜在空间中的扩散过程
在获得低维潜在表示 ( z ) (z) (z)后，LDM 在此空间中执行扩散过程。类似于传统 DMs，前向过程将 ( z ) (z) (z)逐步加噪，后向过程通过一个条件Unet ( ϵ ( z t , t ) ) (\epsilon(z_t,t)) (ϵ(zt,t))预测噪声，逐步生成 ( z ) (z) (z)的样本。最终，生成的 ( z ) (z) (z)通过解码器 ( D ) (\mathcal{D}) (D)转换为图像。
目标函数为
L L D M : = E ϵ ( x ) ， ϵ ∼ N ( 0 , 1 ) , t [ ∣ ∣ ϵ − ϵ t h e t a ( z t , t ) ∣ ∣ 2 2 ] L_{LDM}:=E_{\epsilon(x)，\epsilon \sim N(0,1),t}[||\epsilon-\epsilon_{theta}(z_t,t)||_2^2] LLDM:=Eϵ(x)，ϵ∼N(0,1),t[∣∣ϵ−ϵtheta(zt,t)∣∣22]
python
class Autoencoder(nn.Module):
def __init__(self):
super(Autoencoder, self).__init__()
# 编码器
self.encoder = nn.Sequential(
nn.Conv2d(channels, 16, 4, stride=2, padding=1), # [batch, 16, 14, 14]
nn.ReLU(),
nn.Conv2d(16, latent_dim, 4, stride=2, padding=1), # [batch, latent_dim, 7, 7]
nn.ReLU()
)
# 解码器
self.decoder = nn.Sequential(
nn.ConvTranspose2d(latent_dim, 16, 4, stride=2, padding=1), # [batch, 16, 14, 14]
nn.ReLU(),
nn.ConvTranspose2d(16, channels, 4, stride=2, padding=1), # [batch, 1, 28, 28]
nn.Tanh()
)
def forward(self, x):
z = self.encoder(x)
x_recon = self.decoder(z)
return x_recon, z编解码器通过上卷积和下卷积操作实现。
2 VAE
上面介绍的结构是自编码器的，在自编码器中潜在表示是一个固定值，在变分自编码器(VAE)中，潜在表示是一个不确定变量。
VAE不再只学习提取输入数据的编码信息，而是去学习获取输入数据的概率分布。
与此同时，中间量我们不再叫潜在变量，而是称为潜在空间（也称潜在分布，隐变量等），并使用 Z 来表示，用来凸显其是一个变化的量。
要做一个生成模型，首先我们看生成式模型的"梦想"。我们有一批数据样本 {x_1, ..., x_n}，其所在的集合我们用 X 来表示，我们本想根据 {x_1, ..., x_n} 得到 X 的分布 p(x)，如果能得到的话，那我直接根据 p(x) 来采样，就可以生成所有可能的 x 了（包括 {x_1, ..., x_n} 以外的，这不就是在真正地"生成"新的数据么），这就算是一个终极理想的生成模型了。
当然，这个"梦想"很难实现，于是我们整了一个迂回的策略。具体而言，我们可以借助一个简单的分布 p(z) ，建立其中 z 和 x 的对应关系，也就是 p(x|z)，然后把所有的 z 全都拉过来积分，这样不就能求 p(x) ！如果真的有了这样的一个映射，那我们在推理的时候就可以先从简单分布中采样得到 z，然后走这个映射，然后就可以生成最终的目标 p(x|z) 了！这样的生成方式貌似也不错，整个过程就变成了 z ∼ p ( z ) → x ∼ p ( x ∣ z ) z\sim p(z) → x\sim p(x|z) z∼p(z)→x∼p(x∣z)。于是乎，p(x) 的求法就变成了：
p ( x ) = ∑ z p ( x ∣ z ) p ( z ) = ∫ p ( x ∣ z ) p ( z ) d z p(x)=\sum_z p(x|z)p(z)=\int p(x|z)p(z)dz p(x)=∑zp(x∣z)p(z)=∫p(x∣z)p(z)dz
解码器要做的工作就是如何从这样的概率分布中采样并还原成最终的输出。
p ( X ) = ∫ p ( x ∣ z ) p ( z ) d z p(X)=\int p(x|z)p(z)dz p(X)=∫p(x∣z)p(z)dz
从训练的角度来说，我们需要最大化 p(x)。其中，p(z) 是我们找的一个简单分布，那自然是可知的，那么 p(x|z) 我们怎么求？p(x|z) 是已知潜在表示 z 求对应的 x，如果直接根据公式求 p(x) 的话我们需要列举所有的 z 并做积分，这肯定是不可能的。那这 z 应该从哪里来？好在我们可以用手上已有的数据去得到部分的 z 。具体而言，我们可以用已有的样本数据集 {x_1, ..., x_n} 去得到基于已有样本的 z ，也即后验分布 p(z|x)，从而我们可以根据这些部分 z ，去近似求取 p(x)，或者从训练角度来说，去最大化 p(x)。
实际上，有关"为什么使用这些部分的 z 去最大化 p(x) 是有效的"这个问题，这里还有一种理解方式。我们可以说，对于 z~p(z) 的大部分采样，也就是大部分 z，对于 p(x) 的积分计算都没有贡献，也就是它们的 p(z|x) 都近乎为0，根本就跟 x 没什么关系。我们举个例子。如果定义 Z 是工具集合，目标 X 是职业集合，我自然可以说已知菜刀，厨师的概率会比较大，即 p(厨师|菜刀) 的值是很可观的。这意味着菜刀对于 p(厨师) 的积分计算贡献很大。究其本质原因，厨师本来就经常使用菜刀，即 p(菜刀|厨师) 是大的，这才使得菜刀对厨师分布的计算贡献大。此时，我使用菜刀以及其他厨具去近似计算厨师的分布，这才是有效的，因为他们对厨师的贡献占据了绝大部分。其余的例如什么螺丝刀、雨伞、梯子这类，p(厨师|螺丝刀/雨伞/梯子) 应该很小吧？这当然对最终的厨师分布积分计算没啥用，因此我们为了计算方便，去舍弃这些工具也是合理的，毕竟我们很难列举完世界上所有的工具。当然，这个假设是建立在工具 Z 自己的分布是个均匀分布，即所有先验 p(z) 是相同的，这里也为后续的 ELBO 计算埋下了伏笔。
所以说，再次切换到原来的场景中，我们使用已有的样本数据集 {x_1, ..., x_n} 去得到基于已有样本的 z ，这个过程本身就是在得到 p(z|x) 大的 z ，这些得来的 z 本身就对最终的 p(x) 贡献非常大，和 x 是更相关的，因此我们使用这些从后验分布中得到的 z 去进行训练是有效的、高效的。
讲到现在，我们提到了 p(z|x) 和 p(x|z) ，我们怎么去通过神经网络去实现这些呢？欸！前者是根据 x 求 z，这不就是 Encoder 么？后者是根据 z 求 x，这不就是 Decoder 么？我们很自然就可以引入编解码器架构。首先看 p(z|x)，我们肯定没法计算真正的后验分布，我们只能使用已有的有限的样本数据集 {x_1, ..., x_n} 通过 Encoder 去得到 z，这建模出来的是一个近似后验分布 q，是近似代替真实后验分布 p(z|x)，式子如下：
q ( z ∣ x ) q(z|x) q(z∣x)
好了，如此一来我们就可以根据这个近似的分布 q ，采样对应的 z~q，去走 Decoder ，求取 E_{z~q}[ log p(x|z) ] 了。不过，我们的最终目的是通过训练最大化 p(x) ，那我们应该如何利用这两个式子来最大化 p(x) 呢？我们在这里使用 KL 散度来将这个近似后验分布同我们要求的联系起来。由 KL 散度公式，我们可以有下式 ：
使用贝叶斯公式转化右式的 p(z|x)，式子变形如下：
p(X) 是确定量，我们将其抽出，并进行移项：
再使用 KL 散度公式进行重写，我们就得到了这个式子：
由于 KL 散度非负，因此我们可以说：
上面式子的 RHS（右式）很明显可以算是 log p(x) 的下界，我们一般称其为变分下界（Variational Lower Bound，ELBO）。实际上，我们的目标和 ELBO 的差距就是近似后验分布 q(z|x) 和真实后验分布 p(z|x) 的差距。至此，我们将 "使 p(x) 最大化" 这一问题转化为了将变分下界最大化。
根据上面的理论，我们可以画出我们的模型架构图：
RCNN
这一周还对Selective Search算法进行了实现
python
import numpy as np
import torch
import matplotlib.pyplot as plt
from PIL import Image
from scipy import ndimage
torch.manual_seed(42)
def calculate(data):
L = data[:, :, 0].astype(np.float32) # (H, W)
weight = np.abs(L[:, 1:] - L[:, :-1]) # (H, W-1)
height = np.abs(L[1:, :] - L[:-1, :]) # (H-1, W)
return weight, height
def match(weight, height, data):
H, W = data.shape[:2]
edge = {}
for i in range(H):
for j in range(W - 1):
node1 = i * W + j
node2 = i * W + j + 1
edge[f'({node1},{node2})'] = float(weight[i, j]) # weight is numpy
for i in range(H - 1):
for j in range(W):
node1 = i * W + j
node2 = (i + 1) * W + j
edge[f'({node1},{node2})'] = float(height[i, j])
return edge
class Union:
def __init__(self, H, W):
self.node = list(range(H * W))
self.number = [1] * (H * W)
def find(self, x):
if self.node[x] != x:
self.node[x] = self.find(self.node[x])
return self.node[x]
def union(self, x, y):
rx, ry = self.find(x), self.find(y)
if rx == ry:
return
if self.number[rx] < self.number[ry]:
rx, ry = ry, rx
self.node[ry] = rx
self.number[rx] += self.number[ry]
def segmentation(edge, data, k=0.5, min_size=50):
# data: (H, W, 3) numpy array, uint8
H, W = data.shape[:2]
edge_sort = sorted(edge.items(), key=lambda x: x[1])
union = Union(H, W)
d = [0.0] * (H * W)
for edge_name, weight_val in edge_sort:
nodes = edge_name.strip('()').split(',')
node1 = int(nodes[0])
node2 = int(nodes[1])
area1 = union.find(node1)
area2 = union.find(node2)
if area1 == area2:
continue
m_int1 = d[area1] + k / union.number[area1]
m_int2 = d[area2] + k / union.number[area2]
m_int = min(m_int1, m_int2)
if weight_val <= m_int:
old_int1, old_int2 = d[area1], d[area2]
union.union(node1, node2)
new_area = union.find(node1)
d[new_area] = max(old_int1, old_int2, weight_val)
# 构建 mask
img_mask = np.zeros((H, W), dtype=int)
root_to_label = {}
label_counter = 0
for i in range(H * W):
root = union.find(i)
if root not in root_to_label:
root_to_label[root] = label_counter
label_counter += 1
r, c = divmod(i, W)
img_mask[r, c] = root_to_label[root]
# 小区域合并
labels, counts = np.unique(img_mask, return_counts=True)
small_regions = [label for label, cnt in zip(labels, counts) if cnt < min_size]
if small_regions:
region_mean_color = {}
for label in labels:
mask = (img_mask == label)
if np.any(mask):
region_mean_color[label] = data[mask].mean(axis=0)
else:
region_mean_color[label] = np.array([0, 0, 0])
new_mask = img_mask.copy()
for label in small_regions:
mask = (new_mask == label)
if not np.any(mask):
continue
neighbors = set()
yx = np.argwhere(mask)
for y, x in yx:
for dy, dx in [(-1,0), (1,0), (0,-1), (0,1)]:
ny, nx = y + dy, x + dx
if 0 <= ny < H and 0 <= nx < W:
nb_label = new_mask[ny, nx]
if nb_label != label:
neighbors.add(nb_label)
if not neighbors:
continue
target_color = region_mean_color[label]
best_neighbor = min(neighbors, key=lambda nb: np.linalg.norm(target_color - region_mean_color[nb]))
new_mask[mask] = best_neighbor
img_mask = new_mask
return union, img_mask
def region_features(img_lab, mask, region_id):
# img_lab: (H, W, 3) numpy, uint8 or float
H, W = mask.shape
pixels = img_lab[mask == region_id]
if len(pixels) == 0:
return None
color_hist = []
for ch in range(3):
hist, _ = np.histogram(pixels[:, ch], bins=25, range=(0, 255), density=True)
color_hist.extend(hist)
L = img_lab[:, :, 0].astype(np.float32)
sobel_x = ndimage.sobel(L, axis=0)
sobel_y = ndimage.sobel(L, axis=1)
magnitude = np.hypot(sobel_x, sobel_y)
mag_region = magnitude[mask == region_id]
texture_hist = []
for _ in range(3):
hist, _ = np.histogram(mag_region, bins=10, range=(0, magnitude.max() + 1e-5), density=True)
texture_hist.extend(hist)
yx = np.argwhere(mask == region_id)
min_y, min_x = yx.min(axis=0)
max_y, max_x = yx.max(axis=0)
center_y = (min_y + max_y) / 2
center_x = (min_x + max_x) / 2
return {
'color_hist': np.array(color_hist),
'texture_hist': np.array(texture_hist),
'size': len(pixels),
'bbox': (int(min_y), int(min_x), int(max_y), int(max_x)),
'center': (center_y, center_x)
}
def similarity(r1, r2):
def hist_intersection(h1, h2):
return np.sum(np.minimum(h1, h2))
color = hist_intersection(r1['color_hist'], r2['color_hist'])
texture = hist_intersection(r1['texture_hist'], r2['texture_hist'])
scale = r1['size'] + r2['size'] # 图像总像素数可替换为常数
size = 1.0 - (r1['size'] + r2['size']) / scale
br1 = r1['bbox']
br2 = r2['bbox']
bbr = (
min(br1[0], br2[0]),
min(br1[1], br2[1]),
max(br1[2], br2[2]),
max(br1[3], br2[3])
)
area_bbr = (bbr[2] - bbr[0] + 1) * (bbr[3] - bbr[1] + 1)
area_r1 = (br1[2] - br1[0] + 1) * (br1[3] - br1[1] + 1)
area_r2 = (br2[2] - br2[0] + 1) * (br2[3] - br2[1] + 1)
fill = 1.0 - (area_bbr - area_r1 - area_r2) / area_bbr
return color + texture + size + fill
def selective_search(img_lab, initial_mask):
unique_ids = np.unique(initial_mask)
regions = {}
for rid in unique_ids:
feat = region_features(img_lab, initial_mask, rid)
if feat:
regions[rid] = feat
# 初始化相似度队列
from heapq import heappush, heappop
sim_queue = []
pair_keys = {}
region_ids = list(regions.keys())
for i in range(len(region_ids)):
for j in range(i + 1, len(region_ids)):
r1_id, r2_id = region_ids[i], region_ids[j]
# 只考虑相邻区域（可选：检查 bbox 是否相邻）
if not adjacent(regions[r1_id]['bbox'], regions[r2_id]['bbox']):
continue
s = similarity(regions[r1_id], regions[r2_id])
heappush(sim_queue, (-s, r1_id, r2_id)) # max-heap via negative
pair_keys[(min(r1_id, r2_id), max(r1_id, r2_id))] = True
proposals = [r['bbox'] for r in regions.values()] # 初始区域也作为候选
# 合并过程
while len(sim_queue) > 0 and len(regions) > 1:
_, r1_id, r2_id = heappop(sim_queue)
if r1_id not in regions or r2_id not in regions:
continue # 已被合并
# 合并区域
new_id = max(regions.keys()) + 1
br1, br2 = regions[r1_id]['bbox'], regions[r2_id]['bbox']
new_bbox = (
min(br1[0], br2[0]),
min(br1[1], br2[1]),
max(br1[2], br2[2]),
max(br1[3], br2[3])
)
new_size = regions[r1_id]['size'] + regions[r2_id]['size']
new_center = (
(regions[r1_id]['center'][0] * regions[r1_id]['size'] +
regions[r2_id]['center'][0] * regions[r2_id]['size']) / new_size,
(regions[r1_id]['center'][1] * regions[r1_id]['size'] +
regions[r2_id]['center'][1] * regions[r2_id]['size']) / new_size
)
regions[new_id] = {
'bbox': new_bbox,
'size': new_size,
'center': new_center,
'color_hist': (regions[r1_id]['color_hist'] * regions[r1_id]['size'] +
regions[r2_id]['color_hist'] * regions[r2_id]['size']) / new_size,
'texture_hist': (regions[r1_id]['texture_hist'] * regions[r1_id]['size'] +
regions[r2_id]['texture_hist'] * regions[r2_id]['size']) / new_size
}
del regions[r1_id]
del regions[r2_id]
# 添加新候选框
proposals.append(new_bbox)
for rid in list(regions.keys()):
if rid == new_id:
continue
if adjacent(regions[new_id]['bbox'], regions[rid]['bbox']):
s = similarity(regions[new_id], regions[rid])
heappush(sim_queue, (-s, new_id, rid))
pair_keys[(min(new_id, rid), max(new_id, rid))] = True
return proposals
def adjacent(bbox1, bbox2, threshold=2):
y1_min, x1_min, y1_max, x1_max = bbox1
y2_min, x2_min, y2_max, x2_max = bbox2
if x1_max < x2_min - threshold or x2_max < x1_min - threshold:
return False
if y1_max < y2_min - threshold or y2_max < y1_min - threshold:
return False
return True
if __name__ == '__main__':
image = Image.open('data/rcnn/JPEGImages/000001.jpg').convert('LAB')
weight, height = calculate(np.array(image))
edge = match(weight, height, np.array(image))
union, img_mask = segmentation(edge, np.array(image), k=0.5,min_size=50)
proposals = selective_search(np.array(image), img_mask)
plt.figure(figsize=(10, 8))
plt.imshow(image)
for i, (y1, x1, y2, x2) in enumerate(proposals[:200]):
rect = plt.Rectangle((x1, y1), x2 - x1, y2 - y1,
linewidth=1, edgecolor='r', facecolor='none')
plt.gca().add_patch(rect)
plt.title('Selective Search Proposals')
plt.axis('off')
plt.show()在这个selective search的基础上还实现了一个RCNN模型
python
import os
import xml.etree.ElementTree as ET
import torch
from PIL import Image, ImageDraw
from torch import nn
from torchvision.transforms import transforms
from torch.utils.data import Dataset, DataLoader
from selective_search import *
def compute_iou(box, boxes):
box = torch.as_tensor(box, dtype=torch.float32)
boxes = torch.as_tensor(boxes, dtype=torch.float32)
x1 = torch.max(box[0], boxes[:, 0])
y1 = torch.max(box[1], boxes[:, 1])
x2 = torch.min(box[2], boxes[:, 2])
y2 = torch.min(box[3], boxes[:, 3])
inter_area = torch.clamp(x2 - x1, min=0) * torch.clamp(y2 - y1, min=0)
box_area = (box[2] - box[0]) * (box[3] - box[1])
boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
union_area = box_area + boxes_area - inter_area
iou = inter_area / (union_area + 1e-6)
return iou
class MyDataset(Dataset):
def __init__(self, img_path, label_path, nums=200, transform=None):
self.img_path = img_path
self.label_path = label_path
self.transform = transform
self.nums = nums
self.images = [f for f in os.listdir(self.img_path)]
self.labels = [f for f in os.listdir(self.label_path)]
self.classes = ['__background__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
def __len__(self):
return len(self.images)
def __getitem__(self, id):
img_name = self.images[id]
img_path = os.path.join(self.img_path, img_name)
image = Image.open(img_path).convert('RGB')
weight, height = calculate(np.array(image.convert('LAB')))
edge = match(weight, height, np.array(image.convert('LAB')))
union, img_mask = segmentation(edge, np.array(image.convert('LAB')), k=0.5, min_size=100)
proposals = selective_search(np.array(image.convert('LAB')), img_mask)
proposals = proposals[:self.nums]
bbox, label = self.annotation(id)
bbox = torch.tensor(bbox, dtype=torch.float32)
label = torch.tensor(label, dtype=torch.int64)
crops = []
proposal_labels = []
proposal_boxes = [] # will store [dx, dy, dw, dh]
for (y1, x1, y2, x2) in proposals:
prop_box = [x1, y1, x2, y2] # (x1, y1, x2, y2)
crop = image.crop((x1, y1, x2, y2))
crops.append(self.transform(crop))
ious = compute_iou(prop_box, bbox)
max_iou, idx = ious.max(dim=0)
if max_iou >= 0.5:
l = label[idx].item()
gt_box = bbox[idx] # [x1g, y1g, x2g, y2g]
# Proposal box
px1, py1, px2, py2 = float(x1), float(y1), float(x2), float(y2)
pcx = (px1 + px2) / 2.0
pcy = (py1 + py2) / 2.0
pw = px2 - px1
ph = py2 - py1
# Ground truth box
gx1, gy1, gx2, gy2 = gt_box.tolist()
gcx = (gx1 + gx2) / 2.0
gcy = (gy1 + gy2) / 2.0
gw = gx2 - gx1
gh = gy2 - gy1
# Avoid division by zero
pw = max(pw, 1e-6)
ph = max(ph, 1e-6)
gw = max(gw, 1e-6)
gh = max(gh, 1e-6)
dx = (gcx - pcx) / pw
dy = (gcy - pcy) / ph
dw = torch.log(torch.tensor(gw / pw))
dh = torch.log(torch.tensor(gh / ph))
box_offset = torch.tensor([dx, dy, dw, dh], dtype=torch.float32)
else:
l = 0
box_offset = torch.zeros(4, dtype=torch.float32)
proposal_boxes.append(box_offset)
proposal_labels.append(l)
while len(crops) < self.nums:
crops.append(torch.zeros(3, 64, 64))
proposal_boxes.append(torch.zeros(4))
proposal_labels.append(0)
return torch.stack(crops), torch.tensor(proposal_labels, dtype=torch.long), torch.stack(
proposal_boxes) # shape: (N, 4) with [dx, dy, dw, dh]
def annotation(self, id):
label_name = self.labels[id]
label_path = os.path.join(self.label_path, label_name)
tree = ET.parse(label_path)
root = tree.getroot()
bbox = []
label = []
for obj in root.iter('object'):
bbox.append([
int(obj.find('bndbox/xmin').text),
int(obj.find('bndbox/ymin').text),
int(obj.find('bndbox/xmax').text),
int(obj.find('bndbox/ymax').text)
])
name = obj.find('name').text
label.append(self.class_to_idx[name])
return bbox, label
class RCNN(nn.Module):
def __init__(self, d_model=256, num_classes=21):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(3, d_model, 3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(d_model, 2 * d_model, 3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(2 * d_model, 4 * d_model, 3, stride=1, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d((1, 1))
)
self.fc = nn.Linear(4 * d_model, 8 * d_model)
self.classifier = nn.Sequential(
nn.Linear(8 * d_model, 4 * d_model),
nn.ReLU(),
nn.Linear(4 * d_model, num_classes)
)
self.regressor = nn.Sequential(
nn.Linear(8 * d_model, 4 * d_model),
nn.ReLU(),
nn.Linear(4 * d_model, 4) # bounding box dx, dy, dw, dh
)
def forward(self, x):
x = self.conv(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
logits = self.classifier(x)
bbox = self.regressor(x)
return logits, bbox
def train(model, epoches, train_loader, cls_loss_func, reg_loss_func, optimizer, device):
model.train()
for epoch in range(epoches):
total_cls_loss = 0
total_reg_loss = 0
total_loss = 0
cls_correct = 0
cls_total = 0
for batch_idx, (crops, labels, bbox_offsets) in enumerate(train_loader):
crops = crops.to(device)
labels = labels.to(device)
bbox_offsets = bbox_offsets.to(device)
batch_size, num_proposals = crops.shape[0], crops.shape[1]
crops = crops.view(-1, crops.shape[2], crops.shape[3], crops.shape[4])
labels = labels.view(-1)
bbox_offsets = bbox_offsets.view(-1, 4)
optimizer.zero_grad()
cls_logits, bbox_pred = model(crops)
cls_loss = cls_loss_func(cls_logits, labels)
pos_mask = (labels > 0) # 背景类为0
if pos_mask.sum() > 0:
reg_loss = reg_loss_func(bbox_pred[pos_mask], bbox_offsets[pos_mask])
else:
reg_loss = torch.tensor(0.0, device=device)
loss = cls_loss + reg_loss
loss.backward()
optimizer.step()
total_cls_loss += cls_loss.item()
total_reg_loss += reg_loss.item() if pos_mask.sum() > 0 else 0
total_loss += loss.item()
non_bg_mask = labels >= 0 # 所有样本
if non_bg_mask.sum() > 0:
pred = cls_logits[non_bg_mask].argmax(dim=1)
cls_correct += (pred == labels[non_bg_mask]).sum().item()
cls_total += non_bg_mask.sum().item()
cls_accuracy = 100. * cls_correct / cls_total if cls_total > 0 else 0
print(f'Epoch {epoch + 1}, Loss: {total_loss / len(train_loader):.4f}, Cls_Accuracy: {cls_accuracy:.2f}%')
def load_model(model_path, device):
model = RCNN(d_model=32, num_classes=21).to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
return model
def predict(model, image, transform, device):
with torch.no_grad():
image_tensor = transform(image).unsqueeze(0).to(device)
cls_logits, bbox_pred = model(image_tensor)
pred_label = cls_logits.argmax(dim=1).item()
pred_bbox = bbox_pred.squeeze().cpu().numpy()
return pred_label, pred_bbox
def visualize_predictions(image, predictions, idx_to_class, threshold=0.5):
draw = ImageDraw.Draw(image)
label, bbox = predictions
if label > 0: # 过滤掉背景类别
class_name = idx_to_class[label]
draw.rectangle(bbox, outline="red", width=3)
draw.text((bbox[0], bbox[1]), f'{class_name}', fill="red")
return image
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor()
])
dataset = MyDataset(img_path='data/rcnn/JPEGImages',
label_path='data/rcnn/Annotations',
nums=100,
transform=transform)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
model = RCNN(d_model=32, num_classes=21).to(device)
cls_loss_func = nn.CrossEntropyLoss()
reg_loss_func = nn.SmoothL1Loss() # 常用于边界框回归
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
print('device:', device)
print(f"模型参数量: {sum(p.numel() for p in model.parameters())}")
train(model, epoches=5, train_loader=dataloader, cls_loss_func=cls_loss_func, reg_loss_func=reg_loss_func,
optimizer=optimizer, device=device)
torch.save(model,'data/rcnn/model.pth')
test_image = Image.open('data/rcnn/Test/000301.jpg').convert('RGB')
prediction = predict(model, test_image, transform, device)
idx_to_class = {idx: cls for cls, idx in dataset.class_to_idx.items()}
visualized_image = visualize_predictions(test_image, prediction, idx_to_class)
visualized_image.show()模型由于需要先对图片selective search处理，花费的时间很长，模型的参数量倒是不大，但是在三百张图片的前提下训练都已经很久了。还有Fast-RCNN和其他的改进。batch设置为1是因为不同的图片可能会有不同数目的类别，如果盲目stack会提示维度不匹配。
总结
本周学习了潜在扩散模型和VAE，同时通过代码实现了RCNN，由于手写selective search的原因，代码训练一次的时间很长，所以需要进行优化。