当前位置：首页 > news >正文

别再死记硬背了！用Python和PyTorch从零实现一个Siamese Network（附完整代码）

news 2026/5/31 2:15:41

用Python和PyTorch从零构建孪生网络：实战图像相似度分析

当你第一次听说"孪生网络"时，脑海中浮现的可能是科幻电影里的双胞胎AI。实际上，这种网络结构更像是给计算机安装了一双"火眼金睛"，让它能够辨别两张图片是否属于同一类别。想象一下这样的场景：你手机里有上千张宠物照片，想快速找出所有橘猫的照片；或者电商平台需要自动识别用户上传的商品是否与正品相符。这些正是孪生网络大显身手的领域。

与传统分类网络不同，孪生网络的核心在于比较而非分类。它通过两个共享权重的子网络（因此得名"孪生"）分别处理输入样本，然后比较它们的特征差异。这种设计使其特别适合小样本学习场景——即使每类只有少量样本，也能通过对比学习获得良好的识别效果。下面我们将用PyTorch一步步实现这个神奇的网络，并用常见的猫狗数据集验证其效果。

1. 环境准备与数据加载

工欲善其事，必先利其器。在开始编码前，我们需要配置合适的开发环境。推荐使用Python 3.8+和PyTorch 1.10+版本，这些组合既能保证功能完整又避免最新版本可能存在的兼容性问题。

# 创建虚拟环境（可选但推荐） python -m venv siamese_env source siamese_env/bin/activate # Linux/Mac siamese_env\Scripts\activate # Windows # 安装核心依赖 pip install torch torchvision matplotlib pandas

对于数据集，我们将使用Kaggle经典的"Dogs vs Cats"数据集简化版。这个数据集包含25,000张图片，其中12,500张狗和12,500张猫。为简化实验，我们可以使用预处理后的版本：

import torch from torchvision import datasets, transforms # 定义图像预处理流程 transform = transforms.Compose([ transforms.Resize((100, 100)), # 统一尺寸 transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ImageNet标准化 ]) # 加载数据集 full_dataset = datasets.ImageFolder(root='data/train', transform=transform)

关键细节：图像标准化参数采用ImageNet的均值和标准差，这是计算机视觉领域的常见做法。虽然我们的数据集与ImageNet不同，但这种预处理有助于模型更快收敛。

2. 构建数据对生成器

孪生网络的训练需要特殊的数据格式——样本对（Pairs）或三元组（Triplets）。我们需要自定义一个DataLoader来生成这些结构：

from torch.utils.data import Dataset import random class SiameseDataset(Dataset): def __init__(self, dataset, pairs_per_image=5): self.dataset = dataset self.pairs_per_image = pairs_per_image self.class_indices = self._build_class_indices() def _build_class_indices(self): # 创建类别到索引的映射 class_indices = {} for idx, (_, label) in enumerate(self.dataset): if label not in class_indices: class_indices[label] = [] class_indices[label].append(idx) return class_indices def __len__(self): return len(self.dataset) * self.pairs_per_image def __getitem__(self, index): # 计算原始图像索引和配对类型 img_idx = index // self.pairs_per_image anchor_img, anchor_label = self.dataset[img_idx] # 50%概率选择同类样本，50%选择不同类 if random.random() < 0.5: # 正样本对 pos_indices = self.class_indices[anchor_label] pair_idx = random.choice(pos_indices) while pair_idx == img_idx: # 避免选择相同图像 pair_idx = random.choice(pos_indices) pair_img, _ = self.dataset[pair_idx] target = torch.tensor(1.0, dtype=torch.float32) else: # 负样本对 neg_labels = [l for l in self.class_indices if l != anchor_label] neg_label = random.choice(neg_labels) pair_idx = random.choice(self.class_indices[neg_label]) pair_img, _ = self.dataset[pair_idx] target = torch.tensor(0.0, dtype=torch.float32) return (anchor_img, pair_img), target

提示：在实际项目中，样本对的生成策略会显著影响模型性能。过于简单的负样本（如完全不同类别的图像）会导致模型无法学习细微差异。

数据生成器的使用示例：

from torch.utils.data import DataLoader siamese_data = SiameseDataset(full_dataset) train_loader = DataLoader(siamese_data, batch_size=32, shuffle=True)

3. 设计孪生网络架构

孪生网络的核心在于权重共享——两个输入分支使用相同的网络结构且共享参数。我们先实现基础的CNN特征提取器：

import torch.nn as nn import torch.nn.functional as F class SiameseNetwork(nn.Module): def __init__(self): super(SiameseNetwork, self).__init__() # 共享的特征提取器 self.cnn = nn.Sequential( nn.Conv2d(3, 64, kernel_size=10), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(64, 128, kernel_size=7), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(128, 128, kernel_size=4), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(128, 256, kernel_size=4), nn.ReLU(inplace=True) ) # 相似度计算的全连接层 self.fc = nn.Sequential( nn.Linear(256*6*6, 4096), nn.Sigmoid() ) def forward_one(self, x): x = self.cnn(x) x = x.view(x.size(0), -1) x = self.fc(x) return x def forward(self, input1, input2): output1 = self.forward_one(input1) output2 = self.forward_one(input2) return output1, output2

架构选择解析：

卷积核尺寸依次递减（10→7→4），这是计算机视觉中的常见模式——随着特征图变小，使用更小的卷积核
最后一层不使用池化，保留更多空间信息
全连接层使用Sigmoid激活，将相似度压缩到[0,1]区间

对比损失函数（Contrastive Loss）的实现：

class ContrastiveLoss(nn.Module): def __init__(self, margin=2.0): super(ContrastiveLoss, self).__init__() self.margin = margin def forward(self, output1, output2, label): euclidean_distance = F.pairwise_distance(output1, output2) loss = torch.mean((1-label) * torch.pow(euclidean_distance, 2) + label * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2)) return loss

注意：margin参数控制着正负样本对之间的距离阈值。太小的margin会导致模型难以区分相似样本，太大则可能使训练难以收敛。

4. 训练过程与可视化

有了数据和模型，现在可以开始训练流程。我们将实现一个完整的训练循环，并添加特征可视化功能：

import matplotlib.pyplot as plt from torch.optim import Adam from sklearn.manifold import TSNE def train(model, train_loader, optimizer, criterion, epochs): model.train() for epoch in range(epochs): total_loss = 0 for batch_idx, (data, targets) in enumerate(train_loader): (img1, img2), label = data optimizer.zero_grad() output1, output2 = model(img1, img2) loss = criterion(output1, output2, label) loss.backward() optimizer.step() total_loss += loss.item() if batch_idx % 100 == 0: print(f'Epoch {epoch+1}, Batch {batch_idx}, Current Loss: {loss.item():.4f}') print(f'Epoch {epoch+1}, Average Loss: {total_loss/len(train_loader):.4f}') # 每5个epoch可视化一次特征空间 if (epoch+1) % 5 == 0: visualize_features(model, train_loader.dataset) def visualize_features(model, dataset): model.eval() features = [] labels = [] # 随机选择200个样本进行可视化 indices = random.sample(range(len(dataset)), 200) for idx in indices: (img1, _), label = dataset[idx] with torch.no_grad(): feature = model.forward_one(img1.unsqueeze(0)) features.append(feature.squeeze().numpy()) labels.append(label.item()) # 使用t-SNE降维 tsne = TSNE(n_components=2, perplexity=30) features_2d = tsne.fit_transform(features) # 绘制散点图 plt.figure(figsize=(10,8)) plt.scatter(features_2d[:,0], features_2d[:,1], c=labels, cmap='coolwarm', alpha=0.6) plt.colorbar() plt.title('t-SNE Visualization of Learned Features') plt.show() model.train()

启动训练的完整代码：

# 初始化模型和优化器 model = SiameseNetwork() criterion = ContrastiveLoss() optimizer = Adam(model.parameters(), lr=0.0005) # 开始训练 train(model, train_loader, optimizer, criterion, epochs=20)

训练技巧：

学习率从0.0005开始，如果损失波动较大可适当减小
批量大小(batch size)影响样本对的多样性，32-64是不错的起点
每轮训练后观察特征空间的可视化，确保同类样本逐渐聚集

5. 模型评估与实战应用

训练完成后，我们需要评估模型在实际任务中的表现。不同于传统分类任务的准确率，孪生网络的评估指标有其特殊性：

def evaluate(model, test_loader, threshold=0.5): model.eval() correct = 0 total = 0 with torch.no_grad(): for (img1, img2), labels in test_loader: output1, output2 = model(img1, img2) distances = F.pairwise_distance(output1, output2) predictions = (distances < threshold).float() correct += (predictions == labels).sum().item() total += labels.size(0) accuracy = 100 * correct / total print(f'Test Accuracy: {accuracy:.2f}% (Threshold: {threshold})') return accuracy

在实际部署时，我们可以将模型封装成方便的API：

class SiamesePredictor: def __init__(self, model_path, threshold=0.5): self.model = SiameseNetwork() self.model.load_state_dict(torch.load(model_path)) self.model.eval() self.threshold = threshold self.transform = transforms.Compose([ transforms.Resize((100, 100)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) def predict(self, img1_path, img2_path): img1 = self._load_image(img1_path) img2 = self._load_image(img2_path) with torch.no_grad(): feat1, feat2 = self.model(img1.unsqueeze(0), img2.unsqueeze(0)) distance = F.pairwise_distance(feat1, feat2).item() similarity = 1 - distance return similarity > self.threshold, similarity def _load_image(self, img_path): img = Image.open(img_path).convert('RGB') return self.transform(img)

使用示例：

predictor = SiamesePredictor('best_model.pth') is_same, confidence = predictor.predict('cat1.jpg', 'cat2.jpg') print(f"Same category: {is_same} (Confidence: {confidence:.2%})")

性能优化方向：