PyTorch分割模型的嵌入

Question

我试图从微调的torchvision.models.segmentation.fcn_resnet50模型中得到平均嵌入。 我测试过的一种方法是将分类器与nn.Identity层交换（返回形状张量[batch\u size，feature\u num，h，w]），并在推理中使用torch.mean

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

但我注意到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()可能会导致意外的结果（特别是内存泄漏和分段错误）。我还尝试将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()应用程序问题没有任何联系。（在此设置中，错误不是分段错误，而是来自F.interpolate()的size mismatch，这是可以预期的）。我能做些什么来解决这个问题？我是否需要覆盖模型的forward()方法（如何覆盖？），或者可能存在其他变体

另外，我知道使用分段模型作为嵌入器可能不是最好的主意，但目前正在尝试具体测试这种方法