Pytorch through load_ state_ Dict load weight

1、 If the pre training network is exactly the same as the current network

net = model(input_channels=input_channels, angle_cls=args.angle_bins)
#  Load pre training model 
pretrained_dict = torch.load(" Pre training weight path ")
net.load_state_dict(pretrained_dict) 

2、 If the names of some network layers in the pre training network and the current network are different
Such as , There is a layer of convolution in the pre training network, which is defined as follows ：

self.point_output = nn.Conv2d(input, output, kernel_size, ...)

But the current network has removed this layer . Load weights in the following ways

net = model(input_channels=input_channels, angle_cls=args.angle_bins)
#  Load pre training model 
pretrained_dict = torch.load(" Pre training weight path ")
model_dict = net.state_dict()
state_dict = {
    k:v for k,v in pretrained_dict.items() if k in sgdn_dict.keys()}
model_dict .update(state_dict)
net.load_state_dict(model_dict)

3、 If the name of some network layers in the pre training network and the current network are the same , But the shape is different
For example, there is a network , Pre training with 4 The tunnel RGBD Image as input ; Now when training , Take the depth map of single channel as input , That is, the weight size of the first layer is different . Load weights in the following ways

net = model(input_channels=input_channels, angle_cls=args.angle_bins)
#  Load pre training model 
pretrained_dict = torch.load(" Pre training weight path ")
model_dict = net.state_dict()
state_dict = {
    k:v for k,v in pretrained_dict.items() if k in sgdn_dict.keys() and v.shape == model_dict [k].shape} 
model_dict .update(state_dict)
net.load_state_dict(model_dict)