时间戳

首页 > 解决方案 > 从 PyTorch 分割模型嵌入

python - 从 PyTorch 分割模型嵌入

问题描述

我试图从微调torchvision.models.segmentation.fcn_resnet50模型中获得平均嵌入。我测试过的一种方法是将分类器与nn.Identity层交换（返回形状的张量 [batch_size, feature_num, h, w]）并用于torch.mean推理

model_embedder.classifier = Identity()
outputs = torch.mean(model_embedder(inputs), dim=[2, 3])

但我注意到forward()Torchvision 分割模型的方法如下：

# ...
x = features["out"]
x = self.classifier(x)
x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
result["out"] = x
# ...

并且使用F.interpolate()on 特征可能会导致意想不到的结果（特别是内存泄漏和分段错误）。我也尝试AvgPooling直接将图层添加到模型中

class EmbeddingAvgPooling(nn.Module):
    def __init__(self):
        super(EmbeddingAvgPooling, self).__init__()
        
    def forward(self, x):
        return torch.mean(x, dim=[2, 3])
  
model_embedder.classifier = EmbeddingAvgPooling()

但当然，它与F.interpolate()应用程序问题没有任何关系。（在此设置中，错误不是分段错误，而是size mismatch来自F.interpolate()，这是可以预期的）。我能做些什么来解决这个问题？我是否需要覆盖forward()模型的方法（如何？）或者可能有其他变体？

PS我知道使用分割模型作为嵌入器可能不是最好的想法，但目前正在尝试专门测试这种方法。

标签： pythoncomputer-visionpytorchimage-segmentationtorchvision

解决方案

推荐阅读