https://github.com/facebookresearch/maskrcnn-benchmark/issues/521
ps:如果只需要fintune的话,直接看第二部分,是可以查看model的各个键值的,然后根据需要进行删减。
ps:后面有个乱七八糟的,是我记录的,后期删改
ps:还有个是可以冻结残差网络的训练层的,eg:resnet50,我冻结了最大层,也就是:FREEZE_CONV_BODY_AT = 4
ResNet50FPNStagesTo5 = tuple(
StageSpec(index=i, block_count=c, return_features=r)
for (i, c, r) in ((1, 3, True), (2, 4, True), (3, 6, True), (4, 3, True))
)
FREEZE_CONV_BODY_AT = 1 train resnet only 16 last layers
FREEZE_CONV_BODY_AT = 2 train resnet only 13 last layers
FREEZE_CONV_BODY_AT = 3 train resnet only 9 last layers
FREEZE_CONV_BODY_AT = 4 train resnet only 3 last layers
Steps1) COCO format将自己的数据集按照coco官方形式进行标注。这样会好一点,测试时候直接用coco官方api就可以了。可用labelme等等,再进行转换,还有对应的json文件。official coco format 2) Creating a |
需要的步骤是:
1 -为特定模型创建cfg
2 -使用load_c2_format函数,将为您提供包含模型字段的dict。在那里,可以通过删除字段等来执行您想要的模型手术
3 - 使用 torch.save保存目标,结构应该还是按dict(model=state_dict).保存,并且,请注意扩展名是.pth而不是pkl
4 - 更改 MODEL.WEIGHT指向更改后的.pth文件地址。
前三步分别对应下面的trim_detectron_model.py脚本,注意别忘了第4步说明如下:
pretrained_path:就是从
MODEL_ZOO下载的作者训练好的结果,(后缀.pkl)
save_path:自己修改后用于fintune的模型地址,
请注意扩展名是.pth
cfg:与模型对应的配置文件。
可以看到自己网络里的参数,比如我的是(截取部分,对应删除就好了):
[ # FAST RCNN层,2个FC和分类回归层 'fc1000.bias', 'fc1000.weight', 'fc6.bias', 'fc6.weight', 'fc7.bias', 'fc7.weight', 'bbox_pred.bias', 'bbox_pred.weight', 'cls_score.bias', 'cls_score.weight', # Mask层: 'mask_fcn1.bias', 'mask_fcn1.weight', 'mask_fcn2.bias', 'mask_fcn2.weight', 'mask_fcn3.bias', 'mask_fcn3.weight', 'mask_fcn4.bias', 'mask_fcn4.weight', 'conv5_mask.bias', 'conv5_mask.weight', 'mask_fcn_logits.bias', 'mask_fcn_logits.weight', # RPN层 'rpn.head.conv.bias', 'rpn.head.conv.weight', 'rpn.head.bbox_pred.bias', 'rpn.head.bbox_pred.weight', 'rpn.head.cls_logits.bias', 'rpn.head.cls_logits.weight' # FPN层 'fpn_inner1.bias', 'fpn_inner1.weight', 'fpn_inner2.bias', 'fpn_inner2.weight', 'fpn_inner3.bias', 'fpn_inner3.weight', 'fpn_inner4.bias', 'fpn_inner4.weight', 'fpn_layer1.bias', 'fpn_layer1.weight', 'fpn_layer2.bias','fpn_layer2.weight', 'fpn_layer3.bias', 'fpn_layer3.weight', 'fpn_layer4.bias', 'fpn_layer4.weight', # 残差层 ps:为啥只有bn1和bn2???? 'conv1.bias', 'conv1.weight', 'bn1.bias', 'bn1.weight', 'layer1.0.downsample.0.bias', 'layer1.0.downsample.1.bias', 'layer1.0.downsample.1.weight', 'layer1.0.downsample.0.weight', 'layer1.0.conv1.bias', 'layer1.0.bn1.bias', 'layer1.0.bn1.weight', 'layer1.0.conv1.weight', 'layer1.0.conv2.bias', 'layer1.0.bn2.bias', 'layer1.0.bn2.weight', 'layer1.0.conv2.weight', ................. 'layer4.0.downsample.0.bias', 'layer4.0.downsample.1.bias', 'layer4.0.downsample.1.weight', 'layer4.0.downsample.0.weight', 'layer4.0.conv1.bias', 'layer4.0.bn1.bias', 'layer4.0.bn1.weight', 'layer4.0.conv1.weight', ................. 'layer4.2.conv3.bias', 'layer4.2.bn3.bias', 'layer4.2.bn3.weight', 'layer4.2.conv3.weight', ] '''
trim_detectron_model.py
import os
import torch
import argparse
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.utils.c2_model_loading import load_c2_format
def removekey(d, listofkeys):
r = dict(d)
for key in listofkeys:
print('key: {} is removed'.format(key))
r.pop(key)
return r
parser = argparse.ArgumentParser(description="Trim Detection weights and save in PyTorch format.")
parser.add_argument(
"--pretrained_path",
default="~/.torch/models/_detectron_35858933_12_2017_baselines_e2e_mask_rcnn_R-50-FPN_1x.yaml.01_48_14.DzEQe4wC_output_train_coco_2014_train%3Acoco_2014_valminusminival_generalized_rcnn_model_final.pkl",
help="path to detectron pretrained weight(.pkl)",
type=str,
)
parser.add_argument(
"--save_path",
default="./pretrained_model/mask_rcnn_R-50-FPN_1x_detectron_no_last_layers.pth",
help="path to save the converted model",
type=str,
)
parser.add_argument(
"--cfg",
default="configs/e2e_mask_rcnn_R_50_FPN_1x.yaml",
help="path to config file",
type=str,
)
args = parser.parse_args()
#
DETECTRON_PATH = os.path.expanduser(args.pretrained_path)
print('detectron path: {}'.format(DETECTRON_PATH))
cfg.merge_from_file(args.cfg)
_d = load_c2_format(cfg, DETECTRON_PATH)
print([key for key in _d]) #['model'] 只有这一个键
print([so for so in newdict["model"]]) #或者执行keys = [k for k in _d['model'].keys()]也可以
newdict = _d
#把后面的值从前面移除
newdict['model'] = removekey(_d['model'],
['cls_score.bias', 'cls_score.weight', 'bbox_pred.bias', 'bbox_pred.weight'])
torch.save(newdict, args.save_path)
print('saved to {}.'.format(args.save_path))
# 检查是否按设计输出:
w = torch.load(args.save_path)
print([k for k in w['model'].keys()])
记录一个错误:
w = torch.load("X-101-32x8d.pth")
然而,发生错误:UnicodeDecodeError:'ascii'编解码器无法解码位置2中的字节0xad:序数不在范围内(128)
可以通过使用pickle来解决这个错误:
import pickle
with open("X-101-32x8d.pkl", "rb") as f: w = pickle.load(f, encoding='latin1')
[loaded_state_dict]函数可以检查函数中的哪些键 ,可以在这里进行操作(maskrcnn-benchmark/maskrcnn_benchmark/utils/model_serialization.py )
不会影响repo中的其他部分。大概是
model_dict = model.state_dict()
loaded_state_dict = {k: v for k, v in loaded_state_dict.items() if k in model_dict and "roi_heads" not in k}
model_dict.update(loaded_state_dict)
如果是这样,则无需经过4个步骤。只用编辑 [load_state_dict method]
align_and_update_state_dicts(model_state_dict, loaded_state_dict)
**步骤1**。运行demo.py以测试您的软件包是否已成功安装。同时,您可以保存预先训练的模型,该模型已从pkl格式转换为pth格式,可以通过pytorch模型加载。该模型是演示中COCODEmo对象的数据,因此您可以保存预训练模型的权重,如下所示
pretrained_model_path = "./pretrained_model/maskRCNN.pth"
check_point = {'state': coco_demo.model.state_dict()}
torch.save(check_point, pretrained_model_path)
这样,无需查看[_load_file fun](https://github.com/facebookresearch/maskrcnn-benchmark/blob/f25c6cff92d32d92abe8965d68401004e90c8bee/maskrcnn_benchmark/utils/checkpoint.py#L117).另外还有个保存权重到OrderedDict的方法:(https://github.com/pytorch/pytorch/blob/828cb18fa35c7c132ab920de1c0dc6d859f152d6/torch/nn/modules/module.py#L602)
retrained_model = torch.load(checkpoint_path)['model']
model_dict = self.model.state_dict()
pretrained_dict = {k: v for k, v in pretrained_model.items() if k in model_dict and "roi_heads" not in k}
model_dict.update(pretrained_dict)
**step 2**. 使用configs文件夹中的yaml文件创建fastercnn或maskRCNN模型。执行此操作时,您需要调用cfg.merge_from_list方法来更改[_C.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 81]到数据集中的类数(后台在此实现中不计算)。(https://github.com/facebookresearch/maskrcnn-benchmark/blob/f25c6cff92d32d92abe8965d68401004e90c8bee/maskrcnn_benchmark/config/defaults.py#L182)
**step 3**.需要创建两个类。
第一个类:[torchvision.datasets.coco.CocoDetection] (
maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/coco.py
class COCODataset(torchvision.datasets.coco.CocoDetection):
获取一个图像及其对应的bbox和标签,而不将Bbox和标签捆绑在BoxList对象中。也就是 return img, target, idx
此类与[COCODataset](https://github.com/facebookresearch/maskrcnn-benchmark/blob/f25c6cff92d32d92abe8965d68401004e90c8bee/maskrcnn_benchmark/data/datasets/coco.py#L9). 在每次迭代中,该类生成一个元组(img,bbox,label),其中img由Image.open(...).convert('RGB')生成;
bbox 是 2d numpy 数组, label是1d 数组(mask的话就是掩码信息).
第二类继承自第一类,这个类的主要职责是将一个图像的真实包装在一个BoxList对象中,类似于[line]
img, anno = super(COCODataset, self).__getitem__(idx) |
(https://github.com/facebookresearch/maskrcnn-benchmark/blob/f25c6cff92d32d92abe8965d68401004e90c8bee/maskrcnn_benchmark/data/datasets/coco.py#L43).
**step 4**. 对于pytorch预训练模型,权重在字典中被分类,其中键是layer的名称。因此,您可以排除不存在的layer或者您想要从头开始训练,如下所示:pretrained_model = torch.load(checkpoint_path)['state']
城市景观数据集cityscapes与coco
def clip_weights_from_pretrain_of_coco_to_cityscapes(f, out_file):
""""""
from maskrcnn_benchmark.config.paths_catalog import COCO_CATEGORIES
from maskrcnn_benchmark.config.paths_catalog import CITYSCAPES_FINE_CATEGORIES
coco_cats = COCO_CATEGORIES
cityscapes_cats = CITYSCAPES_FINE_CATEGORIES
coco_cats_to_inds = dict(zip(coco_cats, range(len(coco_cats))))
cityscapes_cats_to_inds = dict(
zip(cityscapes_cats, range(len(cityscapes_cats)))
)
checkpoint = torch.load(f)
m = checkpoint['model']
weight_names = {
"cls_score": "module.roi_heads.box.predictor.cls_score.weight",
"bbox_pred": "module.roi_heads.box.predictor.bbox_pred.weight",
"mask_fcn_logits": "module.roi_heads.mask.predictor.mask_fcn_logits.weight",
}
bias_names = {
"cls_score": "module.roi_heads.box.predictor.cls_score.bias",
"bbox_pred": "module.roi_heads.box.predictor.bbox_pred.bias",
"mask_fcn_logits": "module.roi_heads.mask.predictor.mask_fcn_logits.bias",
}
representation_size = m[weight_names["cls_score"]].size(1)
cls_score = nn.Linear(representation_size, len(cityscapes_cats))
nn.init.normal_(cls_score.weight, std=0.01)
nn.init.constant_(cls_score.bias, 0)
representation_size = m[weight_names["bbox_pred"]].size(1)
class_agnostic = m[weight_names["bbox_pred"]].size(0) != len(coco_cats) * 4
num_bbox_reg_classes = 2 if class_agnostic else len(cityscapes_cats)
bbox_pred = nn.Linear(representation_size, num_bbox_reg_classes * 4)
nn.init.normal_(bbox_pred.weight, std=0.001)
nn.init.constant_(bbox_pred.bias, 0)
dim_reduced = m[weight_names["mask_fcn_logits"]].size(1)
mask_fcn_logits = Conv2d(dim_reduced, len(cityscapes_cats), 1, 1, 0)
nn.init.constant_(mask_fcn_logits.bias, 0)
nn.init.kaiming_normal_(
mask_fcn_logits.weight, mode="fan_out", nonlinearity="relu"
)
def _copy_weight(src_weight, dst_weight):
for ix, cat in enumerate(cityscapes_cats):
if cat not in coco_cats:
continue
jx = coco_cats_to_inds[cat]
dst_weight[ix] = src_weight[jx]
return dst_weight
def _copy_bias(src_bias, dst_bias, class_agnostic=False):
if class_agnostic:
return dst_bias
return _copy_weight(src_bias, dst_bias)
m[weight_names["cls_score"]] = _copy_weight(
m[weight_names["cls_score"]], cls_score.weight
)
m[weight_names["bbox_pred"]] = _copy_weight(
m[weight_names["bbox_pred"]], bbox_pred.weight
)
m[weight_names["mask_fcn_logits"]] = _copy_weight(
m[weight_names["mask_fcn_logits"]], mask_fcn_logits.weight
)
m[bias_names["cls_score"]] = _copy_bias(
m[bias_names["cls_score"]], cls_score.bias
)
m[bias_names["bbox_pred"]] = _copy_bias(
m[bias_names["bbox_pred"]], bbox_pred.bias, class_agnostic
)
m[bias_names["mask_fcn_logits"]] = _copy_bias(
m[bias_names["mask_fcn_logits"]], mask_fcn_logits.bias
)
print("f: {}\nout_file: {}".format(f, out_file))
torch.save(m, out_file)
对最后一层进行手术的两种示例:
Download the model from MODEL_ZOO
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
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)