Pytorch加载预训练模型前n层

时间:2020-06-21 00:29:09   收藏:0   阅读:96
import torch.nn as nn
import torchvision.models as models

class resnet(nn.Module):
def __init__(self):
super(resnet,self).__init__()
self.model = models.resnet18(pretrained=True)
self.encoder = nn.Sequential(*list(self.model.children())[:-1])
def forward(self,x):
y = self.encoder(x)
return y
net = resnet()
评论(0
© 2014 mamicode.com 版权所有 京ICP备13008772号-2  联系我们:gaon5@hotmail.com
迷上了代码!