如何从Pytorch的Fastern-RCNN中删除某些图层?

时间:2020-01-20 02:28:09

标签: python deep-learning pytorch object-detection

目标:我想使用预先训练的Faster-RCNN模型从图像中提取特征。

我尝试过的事情:我使用以下代码来构建模型:

import torchvision.models as models
from PIL import Image
import torchvision.transforms as T
import torch

# download the pretrained fasterrcnn model
model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()
model.cuda()

# remove [2:] layers
modules = list(model.children())[:2]
model_t=torch.nn.Sequential(*modules)

# load image and extract features
img = Image.open('data/person.jpg')
transform = T.Compose([T.ToTensor()])
img_t = transform(img)
batch_t = torch.unsqueeze(img_t, 0).cuda()
ft = model_t(batch_t)

错误:但是出现以下错误:TypeError: conv2d(): argument 'input' (position 1) must be Tensor, not tuple

请帮助!谢谢!

1 个答案:

答案 0 :(得分:1)

print(model.modules)获取图层名称。然后使用以下命令删除图层:

del model.my_layer_name