如何加载自定义模型权重

[baltam]

问题描述：
加载imagenet预训练权重跑没问题。但是当我尝试着加载自己训练好的模型权重时，出现以下错误：

size mismatch for fc.weight: copying a param with shape torch.Size([5, 2048]) from checkpoint, the shape in current model is torch.Size([1000, 2048]).       
        size mismatch for fc.bias: copying a param with shape torch.Size([5]) from checkpoint, the shape in current model is torch.Size([1000]).

imagenet的训练类别有1000类，而我自己训练好的模型只有5个类，并且不包含在imagenet的训练类别中
当维度不一致的时候，我要如何修改代码才能解决这个问题呢，希望dalao能够指点一下~

[mick.yi]

@baltam get_net函数里面改为加载自己的模型就好