Using pre-trained mask-rcnn models on COCO for initialization

[yaoliUoA]

❓ Questions and Help

Hi

I want to use a pre-trained mask-rcnn model on COCO ("catalog://ImageNetPretrained/FAIR/20171220/X-101-32x8d") to fine-tune mask-rcnn on my own dataset. My dataset has only two class (person vs non-person) so I set the ROI_BOX_HEAD. NUM_CLASSES to 2 in the config-file.

However, this gives on error at training:

size mismatch for roi_heads.box.predictor.cls_score.bias: copying a param with shape torch.Size([81]) from checkpoint, the shape in current model is torch.Size([2]).
size mismatch for roi_heads.box.predictor.cls_score.weight: copying a param with shape torch.Size([81, 1024]) from checkpoint, the shape in current model is torch.Size([2, 1024]).
size mismatch for roi_heads.box.predictor.bbox_pred.bias: copying a param with shape torch.Size([324]) from checkpoint, the shape in current model is torch.Size([8]).
size mismatch for roi_heads.box.predictor.bbox_pred.weight: copying a param with shape torch.Size([324, 1024]) from checkpoint, the shape in current model is torch.Size([8, 1024]).
size mismatch for roi_heads.mask.predictor.mask_fcn_logits.bias: copying a param with shape torch.Size([81]) from checkpoint, the shape in current model is torch.Size([2]).
size mismatch for roi_heads.mask.predictor.mask_fcn_logits.weight: copying a param with shape torch.Size([81, 256, 1, 1]) from checkpoint, the shape in current model is torch.Size([2, 256, 1, 1]).

This error is because that box output layer has two classes in the new architecture, however, the coco model used for initialization has 81 classes. Thus the question is, how to use coco model for initialization without initialize the box output layer?

By the way, if I switch to an imagnet model for initialization ("catalog://ImageNetPretrained/FAIR/20171220/X-101-32x8d"), this is no problem at all.

[chengyangfu]

Not sure this is the best solution.

Download the model from MODEL_ZOO
For example

 wget https://download.pytorch.org/models/maskrcnn/e2e_mask_rcnn_X_101_32x8d_FPN_1x.pth
 // Load the model in python
 python
 import torch
 model = torch.load("e2e_mask_rcnn_X_101_32x8d_FPN_1x.pth ")
 // Remove the previous training parameters. 
 del model['iteration']
 del model['scheduler']
 del model['optimizer']
 // Remove the output layers in COCO, these are the mismatched layers you saw.
 //Second stage prediction
 del model["model"]["module.roi_heads.box.predictor.cls_score.weight"]
 del model["model"]["module.roi_heads.box.predictor.cls_score.bias"]
 del model["model"]["module.roi_heads.box.predictor.bbox_pred.weight"]
 del model["model"]["module.roi_heads.box.predictor.bbox_pred.bias"]
 //mask prediction
 del model["model"]["module.roi_heads.mask.predictor.mask_fcn_logits.weight"]
 del model["model"]["module.roi_heads.mask.predictor.mask_fcn_logits.bias"]
 // RPN
 del model["model"]["module.rpn.head.cls_logits.weight"]
 del model["model"]["module.rpn.head.cls_logits.bias"]
 del model["model"]["module.rpn.head.bbox_pred.weight"]
 del model["model"]["module.rpn.head.bbox_pred.bias"]
 //save the model
 torch.save(model, "modified_model.pth")

Then use modified_model.pth in your MODEL.WEIGHT

[zimenglan-sysu-512]

hi @yaoliUoA
u can do some model surgery, like this:

def delete_net_weights_for_finetune(
	model_file, 
	out_file, 
	rpn_final_convs=False, 
	bbox_final_fcs=True, 
	mask_final_conv=True
):
	del_keys = []
	checkpoint = torch.load(model_file)
	print("keys: {}".format(checkpoint.keys()))
	m = checkpoint['model']

	if rpn_final_convs:
		# 'module.rpn.anchor_generator.cell_anchors.0', 
		# 'module.rpn.anchor_generator.cell_anchors.1', 
		# 'module.rpn.anchor_generator.cell_anchors.2', 
		# 'module.rpn.anchor_generator.cell_anchors.3', 
		# 'module.rpn.anchor_generator.cell_anchors.4'
		# 'module.rpn.head.cls_logits.weight', 
		# 'module.rpn.head.cls_logits.bias', 
		# 'module.rpn.head.bbox_pred.weight', 
		# 'module.rpn.head.bbox_pred.bias',
		del_keys.extend([
			k for k in m.keys() if k.find("rpn.anchor_generator") is not -1
		])
		del_keys.extend([
			k for k in m.keys() if k.find("rpn.head.cls_logits") is not -1
		])
		del_keys.extend([
			k for k in m.keys() if k.find("rpn.head.bbox_pred") is not -1
		])

	if bbox_final_fcs:
		# 'module.roi_heads.box.predictor.cls_score.weight', 
		# 'module.roi_heads.box.predictor.cls_score.bias', 
		# 'module.roi_heads.box.predictor.bbox_pred.weight', 
		# 'module.roi_heads.box.predictor.bbox_pred.bias',
		del_keys.extend([
			k for k in m.keys() if k.find(
				"roi_heads.box.predictor.cls_score"
			) is not -1
		])
		del_keys.extend([
			k for k in m.keys() if k.find(
				"roi_heads.box.predictor.bbox_pred"
			) is not -1
		])

	if mask_final_conv:
		# 'module.roi_heads.mask.predictor.mask_fcn_logits.weight', 
		# 'module.roi_heads.mask.predictor.mask_fcn_logits.bias',
		del_keys.extend([
			k for k in m.keys() if k.find(
				"roi_heads.mask.predictor.mask_fcn_logits"
			) is not -1
		])
	
	for k in del_keys:
		print("del k: {}".format(k))
		del m[k]

	# checkpoint['model'] = m
	print("f: {}\nout_file: {}".format(f, out_file))
	recursively_mkdirs(os.path.dirname(out_file))
	torch.save({"model": m}, out_file)

[fmassa]

Hi @yaoliUoA ,

To complement the information given by @chengyangfu and @zimenglan-sysu-512 , have a look at https://github.com/facebookresearch/maskrcnn-benchmark#finetuning-from-detectron-weights-on-custom-datasets and the referenced issue, it has a lot of discussion on how to approach this problem.

In particular, given that you'll be detecting only people, you might want to make the model surgery in the last layer to re-use the parts of the weight that correspond to people.

I'm closing this issue, but let me know if you have further questions

[zacharymostowsky]

Unsure if this will help anyone but the key names are slightly different in newer versions of pytorch. The below worked for me...

# Remove incompatible parameters
del state_dict["roi_heads.box_predictor.bbox_pred.weight"]
del state_dict["roi_heads.box_predictor.cls_score.weight"]
del state_dict["roi_heads.box_predictor.cls_score.bias"]
del state_dict["roi_heads.box_predictor.bbox_pred.bias"] 

[lingbo666]

@zacharymostowsky @chengyangfu ,I also encountered this error. I know the problem may be in checkpoint, key point and is connected with last checkpoint and OUTPUT_DIR. But I can not really learn the specific problem , so I can not learn where I can do this model surgery. I do not find which files I should change. I'll be quite appreciated.