在运行时获取 ResNet 模型全连接层的输入

Question

找到了解决方案，将其作为此问题的答案留在下方:)

项目信息：分类任务 2 类。

我正在尝试为我在运行时放入模型的每张图像获取模型的全连接层的输出。我计划在模型完成训练或测试所有图像后使用它们以使用 UMAP 进行可视化。

型号：

#Load resnet
def get_model():
    model = torchvision.models.resnet50(pretrained=True)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)
    return model

pl模块相关部分：

class classifierModel(pl.LightningModule):
   def __init__(self, model):
     super().__init__()
     self.model = model
     self.learning_rate = 0.0001

def training_step(self, batch, batch_idx):
        x= batch['image']
        y = batch['targets']
        x_hat = self.model(x)
        output = nn.CrossEntropyLoss()
        loss= output(x_hat,y)
        return loss
 
def test_step(self, batch, batch_idx):
        x= batch['image']
        y = batch['targets']
        x_hat = self.model(x)

是否可以通过在pl模块的init中添加一个空列表，然后在x_hat = model(x)执行后添加输出来实现？我怎么知道 x_hat = model(x) 执行后 out_features 不会立即 deleted/discarded ？

Answer 1

x_hat就是这个向量，就是[batch_size, 2048]。因此，只需将您的训练步骤也修改为 return x_hat.

class classifierModel(pl.LightningModule):
   def __init__(self, model):
     super().__init__()
     self.model = model
     self.learning_rate = 0.0001
     self.fc_outputs = []

   def training_step(self, batch, batch_idx):
       x= batch['image']
       y = batch['targets']
       x_hat = self.model(x)
       self.fc_outputs.append(x_hat)
       output = nn.CrossEntropyLoss()
       loss= output(x_hat,y)
       return loss

x_hat 的值不会被删除，除非您在将这些值赋值到别处之前明确调用 del x_hat。在您已经将 x_hat 的值分配给另一个变量的情况下（在您的情况下，听起来您想将其附加到列表中）与这些值关联的内存地址不会被释放，因为仍然有一个即使在引用这些地址的原始变量之后引用这些地址的变量（x_hat 可能已被删除）。这样，python 在内存引用方面相对安全，因为它会在运行时动态计算何时不再需要内存地址/值。

Answer 2

我能够使用 avgpool 层上的前向挂钩并在每个 test_step 上保存输出，如 here 所述：

#Define Hook: 
def get_features(name):
    def hook(model, input, output):
        features[name] = output.detach()
    return hook

现在，当我加载我的模型时，我注册了钩子：

#Load resnet model:
def get_model():
    model = models.resnet50(pretrained=True)
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)
    model.avgpool.register_forward_hook(get_features('feats')) #register the hook
    return model

我不需要更改pytorch闪电模型的初始化，但测试步骤功能：

FEATS = []
# placeholder for batch features
features = {}

class classifierModel(pl.LightningModule):
   def __init__(self, model):
     super().__init__()
     self.model = model
     self.learning_rate = 0.0001

def test_step(self, batch,batch_idx):
        x= batch['image']
        y = batch['targets']
        x_hat = self.model(x)
        FEATS.append(features['feats'].cpu().numpy()) #added this line to save output

现在我们得到了输出 FEATS[0].shape --> (16, 2048, 1, 1)，这是我想要得到的（16 是使用的批大小）。

在运行时获取 ResNet 模型全连接层的输入

Get input of fully connected layer of ResNet model during runtime

machine-learning

deep-learning

pytorch

pytorch-lightning