mxnet學習(9):使用gluon接口讀取symbol預訓練模型finetune

使用symbol得到的模型或者gluon的hybridize之後的模型包括一個.json文件(網絡結構)和.params文件(參數),gluon可以使用net = gluon.SymbolBlock.imports(json, ['data'], params, ctx)導入網絡和參數,這樣可以進行測試或者進一步訓練。

但是如果只需要使用模型的其中一部分,比如只需要conv層,去掉所有fc層,或者再另外增加一些層, 這樣直接導入就會比較複雜。正確的做法如下:

sym, arg_params, aux_params = mx.model.load_checkpoint("1.0.3", 40)#這裏是model的名字和參數對應的epoch
layers = sym.get_internals()#得到所有的layers
outputs = layers['stage4_unit1_conv2_output']#選擇輸出層
inputs = layers['data']#選擇輸入層
net = gluon.SymbolBlock(outputs, inputs)#使用gluon的接口將其封裝成一個新的net
net.load_parameters("1.0.3-0040.params", ignore_extra = True, allow_missing = True)#載入數據
y = net(data)
print(y.shape)

如果需要在該網絡的基礎上再新增加一些層,如下定義:

class PretrainedNetwork(gluon.HybridBlock):
    def __init__(self, pretrained_layer, **kwargs):
        super(PretrainedNetwork, self).__init__(**kwargs)
        with self.name_scope():
            self.pretrained_layer = pretrained_layer #(n, 4, 4, 128)
            self.fc = nn.HybridSequential()
            self.fc.add(
                        nn.Flatten(),
                        nn.Dense(256, activation = 'relu'),
                        nn.Dropout(rate = 0.5),
                        nn.Dense(128)
                        )
            self.single_fc = nn.Dense(2)
            self.fusion_fc = nn.Dense(2)
            
    def hybrid_forward(self, F, x):
        x = self.pretrained_layer(x)
        x = self.fc(x)
        feat = x
        y1 = self.single_fc(x)
        feat = feat.sum(axis = 1)
        y2 = self.fusion_fc(feat)
        return y1, y2

那麼可以通過下面的方式,使用預訓練模型初始化其中一部分:

net = PretrainedNetwork(pretrained_layer = net)
net.initialize(forece_reinit = False, init = init.Xavier())

需要注意的是,要先load_parameters再用其初始化PretrainedNwtwork,否則容易出現prefix不匹配的問題。

如果需要fix其中一部分參數,只訓練其中一部分,可以通過觀察所有layer的名字,找到需要訓練的layer。

print(net.collect_params())#打印所有的參數,這樣可以看到所有的layer及其參數

Trainerparams通過正則表達式選擇需要訓練的參數:

trainer = gluon.Trainer(params = net.collect_params("pretrained*|dense0*"), optimizer = optimizer)

這樣沒有被選中的參數就會被fix,訓練中不會改變。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章