yixiong/FoodClassifier
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

增加resnet_embedding的代码

Browse Source
This commit is contained in:
zhanghuan
2025-09-17 14:44:34 +08:00
parent 44031b92cd
commit 8818643025
1 changed files with 448 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
net/resnet_embedding.py
+448
View File
@@ -0,0 +1,448 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
from PIL import Image
import numpy as np
from typing import Union, List, Tuple, Optional
from settings import settings
class ResNet50EmbeddingNet(nn.Module):
"""
基于ResNet50的食物特征提取网络
移除分类层，输出512维特征向量用于相似度计算
"""
def __init__(self, embedding_dim: int = 512, pretrained: bool = True, use_internal_preprocess: bool = False):
"""
初始化ResNet50 Embedding网络
Args:
embedding_dim: 输出特征向量维度，默认512
pretrained: 是否使用预训练权重
use_internal_preprocess: 是否在模型内部进行预处理
"""
super(ResNet50EmbeddingNet, self).__init__()
self.embedding_dim = embedding_dim
self.use_internal_preprocess = use_internal_preprocess
# 加载预训练的ResNet50
self.backbone = models.resnet50(pretrained=pretrained)
# 获取ResNet50最后一层的输入特征数（2048）
backbone_output_dim = self.backbone.fc.in_features
# 移除原始的分类层
self.backbone.fc = nn.Identity()
# 添加embedding层
self.embedding_layer = nn.Sequential(
nn.Linear(backbone_output_dim, embedding_dim),
nn.BatchNorm1d(embedding_dim),
nn.ReLU(inplace=True),
nn.Dropout(0.2)
)
# 图片预处理变换（用于推理）
self.preprocess = transforms.Compose([
transforms.Resize((224, 224), interpolation=transforms.InterpolationMode.BILINEAR),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 内部预处理变换（用于已经是tensor但未归一化的数据）
self.internal_preprocess = transforms.Compose([
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def preprocess_image(self, image: Union[Image.Image, np.ndarray]) -> torch.Tensor:
"""
预处理单张图片
Args:
image: PIL Image 或 numpy array
Returns:
torch.Tensor: 预处理后的张量，形状为 (1, 3, 224, 224)
"""
if not isinstance(image, Image.Image):
# 如果是numpy array，转换为PIL Image
if hasattr(image, 'shape'):
if len(image.shape) == 3 and image.shape[2] == 3:
image = Image.fromarray(image.astype(np.uint8))
else:
raise ValueError("numpy array必须是3通道RGB格式")
else:
raise ValueError("输入必须是PIL Image或numpy array")
# 应用预处理变换
tensor = self.preprocess(image)
# 添加batch维度
tensor = tensor.unsqueeze(0)
return tensor
def mobile_preprocess(self, x: torch.Tensor) -> torch.Tensor:
"""
移动端预处理（TorchScript兼容）
Args:
x: 输入tensor，形状为 [batch, 3, height, width]，值范围 0-1
Returns:
torch.Tensor: 预处理后的tensor，形状为 [batch, 3, 224, 224]
"""
# 缩放到 224x224
x = F.interpolate(x, size=(224, 224), mode='bilinear', align_corners=False)
# ImageNet标准化
mean = torch.tensor([0.485, 0.456, 0.406], device=x.device).view(1, 3, 1, 1)
std = torch.tensor([0.229, 0.224, 0.225], device=x.device).view(1, 3, 1, 1)
x = (x - mean) / std
return x
def _forward_backbone(self, x: torch.Tensor) -> torch.Tensor:
"""
通过ResNet50骨干网络提取特征
Args:
x: 预处理后的输入tensor
Returns:
torch.Tensor: 骨干网络输出特征
"""
return self.backbone(x)
def _forward_embedding(self, backbone_features: torch.Tensor) -> torch.Tensor:
"""
将骨干网络特征转换为embedding向量
Args:
backbone_features: 骨干网络输出特征
Returns:
torch.Tensor: embedding向量
"""
return self.embedding_layer(backbone_features)
def forward(self, x: torch.Tensor, normalize: bool = True) -> torch.Tensor:
"""
前向传播
Args:
x: 输入tensor
normalize: 是否对输出进行L2归一化
Returns:
torch.Tensor: embedding向量
"""
# 如果启用内部预处理且输入是tensor
if self.use_internal_preprocess and isinstance(x, torch.Tensor):
x = self.internal_preprocess(x)
# 通过骨干网络
backbone_features = self._forward_backbone(x)
# 生成embedding
embeddings = self._forward_embedding(backbone_features)
# L2归一化（用于余弦相似度计算）
if normalize:
embeddings = F.normalize(embeddings, p=2, dim=1)
return embeddings
def forward_mobile(self, x: torch.Tensor, normalize: bool = True) -> torch.Tensor:
"""
移动端前向传播（包含预处理）
Args:
x: 输入tensor，形状为 [batch, 3, height, width]，值范围 0-1
normalize: 是否对输出进行L2归一化
Returns:
torch.Tensor: embedding向量
"""
# 移动端预处理
x = self.mobile_preprocess(x)
# 执行前向传播
return self.forward(x, normalize=normalize)
def extract_embedding(self, image: Union[Image.Image, np.ndarray, torch.Tensor],
normalize: bool = True) -> np.ndarray:
"""
从单张图片提取embedding向量
Args:
image: 输入图片（PIL Image、numpy array或torch.Tensor）
normalize: 是否进行L2归一化
Returns:
np.ndarray: embedding向量
"""
self.eval()
with torch.no_grad():
if isinstance(image, torch.Tensor):
if len(image.shape) == 3:
image = image.unsqueeze(0) # 添加batch维度
x = image
else:
x = self.preprocess_image(image)
# 移动到正确的设备
device = next(self.parameters()).device
x = x.to(device)
# 提取特征
embedding = self.forward(x, normalize=normalize)
return embedding.cpu().numpy().flatten()
def extract_batch_embeddings(self, images: List[Union[Image.Image, np.ndarray]],
normalize: bool = True, batch_size: int = 32) -> np.ndarray:
"""
批量提取embedding向量
Args:
images: 图片列表
normalize: 是否进行L2归一化
batch_size: 批处理大小
Returns:
np.ndarray: embedding向量数组，形状为 (N, embedding_dim)
"""
self.eval()
embeddings = []
with torch.no_grad():
for i in range(0, len(images), batch_size):
batch_images = images[i:i + batch_size]
# 预处理批次图片
batch_tensors = []
for img in batch_images:
tensor = self.preprocess_image(img)
batch_tensors.append(tensor.squeeze(0)) # 移除batch维度
# 堆叠成批次
batch_tensor = torch.stack(batch_tensors)
# 移动到正确的设备
device = next(self.parameters()).device
batch_tensor = batch_tensor.to(device)
# 提取特征
batch_embeddings = self.forward(batch_tensor, normalize=normalize)
embeddings.append(batch_embeddings.cpu().numpy())
return np.vstack(embeddings)
@staticmethod
def cosine_similarity(embedding1: np.ndarray, embedding2: np.ndarray) -> float:
"""
计算两个embedding向量的余弦相似度
Args:
embedding1: 第一个embedding向量
embedding2: 第二个embedding向量
Returns:
float: 余弦相似度值 [-1, 1]
"""
# 确保向量是一维的
if len(embedding1.shape) > 1:
embedding1 = embedding1.flatten()
if len(embedding2.shape) > 1:
embedding2 = embedding2.flatten()
# 计算余弦相似度
dot_product = np.dot(embedding1, embedding2)
norm1 = np.linalg.norm(embedding1)
norm2 = np.linalg.norm(embedding2)
if norm1 == 0 or norm2 == 0:
return 0.0
return dot_product / (norm1 * norm2)
@staticmethod
def batch_cosine_similarity(query_embedding: np.ndarray,
database_embeddings: np.ndarray) -> np.ndarray:
"""
计算查询向量与数据库中所有向量的余弦相似度
Args:
query_embedding: 查询向量，形状为 (embedding_dim,)
database_embeddings: 数据库向量，形状为 (N, embedding_dim)
Returns:
np.ndarray: 相似度数组，形状为 (N,)
"""
# 确保查询向量是一维的
if len(query_embedding.shape) > 1:
query_embedding = query_embedding.flatten()
# 归一化向量（如果还没有归一化）
query_norm = np.linalg.norm(query_embedding)
if query_norm > 0:
query_embedding = query_embedding / query_norm
database_norms = np.linalg.norm(database_embeddings, axis=1)
database_norms[database_norms == 0] = 1 # 避免除零
database_embeddings_normalized = database_embeddings / database_norms[:, np.newaxis]
# 计算余弦相似度
similarities = np.dot(database_embeddings_normalized, query_embedding)
return similarities
def get_embedding_dim(self) -> int:
"""
获取embedding向量维度
Returns:
int: embedding维度
"""
return self.embedding_dim
def freeze_backbone(self):
"""
冻结骨干网络参数，只训练embedding层
"""
for param in self.backbone.parameters():
param.requires_grad = False
def unfreeze_backbone(self):
"""
解冻骨干网络参数，允许端到端训练
"""
for param in self.backbone.parameters():
param.requires_grad = True
def get_trainable_parameters(self):
"""
获取可训练参数
Returns:
generator: 可训练参数生成器
"""
return filter(lambda p: p.requires_grad, self.parameters())
def create_resnet50_embedding(embedding_dim: int = 512,
pretrained: bool = True,
use_internal_preprocess: bool = False) -> ResNet50EmbeddingNet:
"""
创建ResNet50 Embedding模型
Args:
embedding_dim: embedding向量维度
pretrained: 是否使用预训练权重
use_internal_preprocess: 是否在forward中进行预处理
Returns:
ResNet50EmbeddingNet: 网络模型实例
"""
return ResNet50EmbeddingNet(
embedding_dim=embedding_dim,
pretrained=pretrained,
use_internal_preprocess=use_internal_preprocess
)
def create_mobile_resnet50_embedding(embedding_dim: int = 512) -> ResNet50EmbeddingNet:
"""
创建移动端ResNet50 Embedding模型
Args:
embedding_dim: embedding向量维度
Returns:
ResNet50EmbeddingNet: 配置为移动端使用的网络模型实例
"""
class MobileResNet50Embedding(ResNet50EmbeddingNet):
"""
移动端ResNet50 Embedding模型
重写forward方法以包含预处理
"""
def __init__(self, embedding_dim: int = 512):
super().__init__(embedding_dim=embedding_dim, pretrained=True, use_internal_preprocess=False)
def forward(self, x: torch.Tensor, normalize: bool = True) -> torch.Tensor:
"""
移动端前向传播（自动包含预处理）
Args:
x: 输入tensor，形状为 [batch, 3, height, width]，值范围 0-1
normalize: 是否对输出进行L2归一化
Returns:
torch.Tensor: embedding向量
"""
return self.forward_mobile(x, normalize=normalize)
return MobileResNet50Embedding(embedding_dim=embedding_dim)
if __name__ == "__main__":
# 测试网络
print("=== ResNet50 Embedding网络测试 ===")
# 创建模型
model = create_resnet50_embedding(embedding_dim=512, pretrained=True)
print(f"模型参数数量: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
print(f"Embedding维度: {model.get_embedding_dim()}")
# 测试前向传播
dummy_input = torch.randn(2, 3, 224, 224)
embeddings = model(dummy_input)
print(f"输入形状: {dummy_input.shape}")
print(f"输出embedding形状: {embeddings.shape}")
print(f"输出数值范围: [{embeddings.min():.3f}, {embeddings.max():.3f}]")
# 测试L2归一化
embedding_norms = torch.norm(embeddings, p=2, dim=1)
print(f"L2归一化后的向量模长: {embedding_norms}")
# 测试图片预处理
try:
import numpy as np
# 创建一个测试图片 (RGB格式)
test_image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
# 测试单张图片特征提取
embedding = model.extract_embedding(test_image)
print(f"单张图片embedding形状: {embedding.shape}")
print(f"单张图片embedding模长: {np.linalg.norm(embedding):.3f}")
# 测试批量特征提取
test_images = [test_image] * 3
batch_embeddings = model.extract_batch_embeddings(test_images, batch_size=2)
print(f"批量embedding形状: {batch_embeddings.shape}")
# 测试相似度计算
similarity = ResNet50EmbeddingNet.cosine_similarity(embedding, batch_embeddings[0])
print(f"自相似度: {similarity:.3f}")
# 测试批量相似度计算
similarities = ResNet50EmbeddingNet.batch_cosine_similarity(embedding, batch_embeddings)
print(f"批量相似度: {similarities}")
except Exception as e:
print(f"图片处理测试失败: {e}")
# # 测试移动端模型
# print("\n=== 移动端模型测试 ===")
# mobile_model = create_mobile_resnet50_embedding(embedding_dim=512)
# mobile_input = torch.rand(1, 3, 256, 256) # 模拟移动端输入
# mobile_embedding = mobile_model(mobile_input)
# print(f"移动端输入形状: {mobile_input.shape}")
# print(f"移动端输出embedding形状: {mobile_embedding.shape}")
# print(f"移动端embedding模长: {torch.norm(mobile_embedding, p=2, dim=1)}")
#
# print("\n=== 测试完成 ===")