Pytorch-修改模型权重的大小

cfg='models/yolov5s.yaml'
device = torch_utils.select_device('0,1')
weights = 'weights/qr_model_320.pth'

model = Model(cfg).to(device)
dict = torch.load(weights)
model.load_state_dict(dict,strict=False)
print(type(dict)) #
#print(dict['model.0.conv.conv.weight']) #报错,没有添加module
#print(dict['module.model.0.conv.conv.weight'].shape) #torch.Size([32, 12, 3, 3])

##########################修改部分#############################
dict['module.model.0.conv.conv.weight'].resize_(32, 4, 3, 3)
print(dict['module.model.0.conv.conv.weight'].shape) 
torch.save(dict, 'weights/qr_model_320_gray.pth')

 

你可能感兴趣的:(Pytorch)