如何使用Pytorch冻结你不想要其参数被改变的层

如何使用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()
 
# 优化、正则项、权重设置与冻结层