使用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及其參數
在Trainer
的params
通過正則表達式選擇需要訓練的參數:
trainer = gluon.Trainer(params = net.collect_params("pretrained*|dense0*"), optimizer = optimizer)
這樣沒有被選中的參數就會被fix,訓練中不會改變。