使用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,训练中不会改变。