做了一个向量相乘的APP,避免embedding_mobile递归循环.
This commit is contained in:
File diff suppressed because it is too large
Load Diff
+22
-16
@@ -97,21 +97,31 @@ class ResNet50EmbeddingNet(nn.Module):
|
||||
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]
|
||||
"""
|
||||
# # 通过 scale_factor 避免对 size 的整数检查(规避 RecursionError)
|
||||
# h, w = x.shape[2], x.shape[3]
|
||||
# # 防御:避免除零
|
||||
# if h == 0 or w == 0:
|
||||
# raise ValueError(f"Invalid input size: height={h}, width={w}")
|
||||
# scale_h = 224.0 / float(h)
|
||||
# scale_w = 224.0 / float(w)
|
||||
#
|
||||
# x = F.interpolate(x, scale_factor=(scale_h, scale_w),
|
||||
# mode='bilinear', align_corners=False)
|
||||
# 缩放到 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:
|
||||
@@ -180,8 +190,8 @@ class ResNet50EmbeddingNet(nn.Module):
|
||||
# 移动端预处理
|
||||
x = self.mobile_preprocess(x)
|
||||
|
||||
# 执行前向传播
|
||||
return self.forward(x, normalize=normalize)
|
||||
# 显式调用基类 forward,避免子类重写的 forward 形成递归
|
||||
return ResNet50EmbeddingNet.forward(self, x, normalize=normalize)
|
||||
|
||||
def extract_embedding(self, image: Union[Image.Image, np.ndarray, torch.Tensor],
|
||||
normalize: bool = True) -> np.ndarray:
|
||||
@@ -210,7 +220,8 @@ class ResNet50EmbeddingNet(nn.Module):
|
||||
|
||||
# 提取特征
|
||||
embedding = self.forward(x, normalize=normalize)
|
||||
|
||||
# embedding = self.forward(x, normalize=False)
|
||||
|
||||
return embedding.cpu().numpy().flatten()
|
||||
|
||||
def extract_batch_embeddings(self, images: List[Union[Image.Image, np.ndarray]],
|
||||
@@ -365,7 +376,7 @@ def create_resnet50_embedding(embedding_dim: int = 512,
|
||||
)
|
||||
|
||||
|
||||
def create_mobile_resnet50_embedding(embedding_dim: int = 512) -> ResNet50EmbeddingNet:
|
||||
def create_mobile_resnet50_embedding(embedding_dim: int = 512,pretrained: bool = True) -> ResNet50EmbeddingNet:
|
||||
"""
|
||||
创建移动端ResNet50 Embedding模型
|
||||
|
||||
@@ -387,15 +398,10 @@ def create_mobile_resnet50_embedding(embedding_dim: int = 512) -> ResNet50Embedd
|
||||
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)
|
||||
# 在子类里做预处理,然后显式调用基类 forward
|
||||
x = self.mobile_preprocess(x)
|
||||
return ResNet50EmbeddingNet.forward(self, x, normalize=normalize)
|
||||
|
||||
return MobileResNet50Embedding(embedding_dim=embedding_dim)
|
||||
|
||||
|
||||
@@ -15,7 +15,8 @@ def main():
|
||||
print("开始转换ResNet50 Embedding模型为移动端格式...")
|
||||
|
||||
# 1. 加载训练好的embedding模型权重
|
||||
base_model = create_resnet50_embedding(embedding_dim=512, pretrained=True)
|
||||
# base_model = create_resnet50_embedding(embedding_dim=512, pretrained=True)
|
||||
base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True)
|
||||
model_path = "../model/embedding_20250917_145342/best_embedding_model.pth"
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
|
||||
@@ -127,11 +127,11 @@ class TwoLayerSimilarityNet(nn.Module):
|
||||
raise ValueError("输入张量形状必须为 [batch, 512]")
|
||||
|
||||
# 输入归一化
|
||||
# x_norm = F.normalize(x, p=2, dim=1) # [B, 512]
|
||||
x_norm = F.normalize(x, p=2, dim=1) # [B, 512]
|
||||
|
||||
# 单层线性变换
|
||||
# sims = self.layer2(x_norm) # [B, N]
|
||||
sims = self.layer2(x) # [B, N]
|
||||
sims = self.layer2(x_norm) # [B, N]
|
||||
# sims = self.layer2(x) # [B, N]
|
||||
|
||||
return sims
|
||||
|
||||
|
||||
Reference in New Issue
Block a user