1.前要
题主在使用ResNet的pytorch版本时,首先用了第三方实现版,发现其运行时GPU占用率特别高,这是不正常的,最后果断想到使用torchversion中的官方实现版,发现其GPU占用率果然比第三方优化很多,batch可以提高很多。但是官方实现版是1000类的输出,我们需要自己定义的输出类型,比如17。这样的话,就没办法加载预训练模型,所以我们想到以下的解决办法,可以成功实现预训练模型的加载。
2实现方法
import torch
import torch.nn as nn
from torchvision.models.resnet import resnet50
class Net(nn.Module)