How to load a pre-trained model that is with customized loss

I have trained resnet50 model with mulit-task losses, and each loss is customized so the name is not softmax, but makeloss_XX.
I want to use this model and pre-trained model for other tasks, which has different losses, then I got errors. Seems when I loaded it load every layer including the makeloss_XX layer, once the new loss is with different structure, it reports error. The error reports the shape doesn’t match.

Question: How do I load this model by removing this layer?

I use this code to load pre-trained model:
sym_tmp, arg_params, aux_params = mx.model.load_checkpoint(base_params_path, ‘2’)
2 is epoch number.


If you want to stick to the Module API you can have a look at this tutorial:

The part that really interests you is the following snippet:

all_layers = sym_tmp.get_internals()
new_sym = all_layers[layer_name+'_output']

where layer_ name is the name of the layer before the loss.