如何使用Pytorch冻结你不想要其参数被改变的层
我猜下面的内容应该是你想要的!
#---------------------------model define ------------------------------------#
class YourModel(nn.Module):
def __init__(self, feat_dim): # input the dim of output fea-map of Resnet:
super(YourModel, self).__init__()
BackBone = models.resnet50(pretrained=True)
add_block = []
add_block += [nn.Linear(2048, 512)]
add_block += [nn.LeakyReLU(inplace=True)]
add_block = nn.Sequential(*add_block)
add_block.apply(weights_init_xavier)
self.BackBone = BackBone
self.add_block = add_block
def forward(self, input): # input is 2048!
x = self.BackBone(input)
x = self.add_block(x)
return x
#----------------------------------------------------------------------------#
# 模型准备
model = YourModel()
# 优化、正则项、权重设置与冻结层