def __init__(self, classes, step2_model_path, step3_model_path, num_layers=101, fix_cnn_base=False, \
class_agnostic=False, pretrained=False, base_model='resnet101'):
super(resnet_step4, self).__init__()
self.step2_model_path = step2_model_path
self.step3_model_path = step3_model_path
self.dout_base_model = 1024
self.pretrained = pretrained
self.class_agnostic = class_agnostic
self.num_layer = num_layers
self.fix_cnn_base = fix_cnn_base
self.classes = classes
self.n_classes = len(classes)
self.base_model = 'resnet'+str(self.num_layer)
# loss
self.RCNN_loss_cls = 0
self.RCNN_loss_bbox = 0
#define base
if self.num_layer == 101:
resnet = resnet101()
elif self.num_layer == 50:
resnet = resnet50()
elif self.num_layer == 152:
resnet = resnet152()
print("Step4: Loading pretrained weights from %s and %s" %(step2_model_path,step3_model_path))
state_dict_step2 = torch.load(self.step2_model_path)
state_dict_step3 = torch.load(self.step3_model_path)
# not using the last maxpool layer
# Build resnet.
self.RCNN_base = nn.Sequential(resnet.conv1, resnet.bn1,resnet.relu,
resnet.maxpool,resnet.layer1,resnet.layer2,resnet.layer3)
self.RCNN_base.load_state_dict({k.replace('RCNN_base.',''):v for k,v in state_dict_step3['model'].items() if 'RCNN_base' in k})
for key, value in dict(self.RCNN_base.named_parameters()).items():
value.requires_grad = False
# define rpn
self.RCNN_rpn = _RPN(self.dout_base_model)
#init weight of rpn
self.RCNN_rpn.load_state_dict({k.replace('RCNN_rpn.',''):v for k,v in state_dict_step3['model'].items() if 'RCNN_rpn' in k})
for key, value in dict(self.RCNN_rpn.named_parameters()).items():
value.requires_grad = False
# define detector
self.detector = _detector(self.classes, self.class_agnostic,pretrained, base_model=base_model)
self.detector.load_state_dict({k.replace('detector.',''):v for k,v in state_dict_step2['model'].items() if 'detector' in k})
def _init_modules(self):
if self.num_layer == 101:
resnet = resnet101()
elif self.num_layer == 50:
resnet = resnet50()
elif self.num_layer == 152:
resnet = resnet152()
if self.pretrained == True:
print("Loading pretrained weights from %s" % (self.model_path))
state_dict = torch.load(self.model_path)
resnet.load_state_dict({
k: v
for k, v in state_dict.items() if k in resnet.state_dict()
})
# Build resnet.
self.RCNN_base = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.maxpool, resnet.layer1,
resnet.layer2, resnet.layer3)
# Fix blocks
for p in self.RCNN_base[0].parameters():
p.requires_grad = False
for p in self.RCNN_base[1].parameters():
p.requires_grad = False
assert (0 <= cfg.RESNET.FIXED_BLOCKS < 4)
if cfg.RESNET.FIXED_BLOCKS >= 3:
for p in self.RCNN_base[6].parameters():
p.requires_grad = False
if cfg.RESNET.FIXED_BLOCKS >= 2:
for p in self.RCNN_base[5].parameters():
p.requires_grad = False
if cfg.RESNET.FIXED_BLOCKS >= 1:
for p in self.RCNN_base[4].parameters():
p.requires_grad = False
def set_bn_fix(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
for p in m.parameters():
p.requires_grad = False
self.RCNN_base.apply(set_bn_fix)
def __init__(self, classes, num_layers=101, fix_cnn_base=False, \
class_agnostic=False, pretrained=False, base_model='resnet101'):
super(resnet_step1, self).__init__()
if num_layers == 101:
self.model_path = 'data/pretrained_model/resnet101_caffe.pth'
elif num_layers == 50:
self.model_path = 'data/pretrained_model/resnet50_caffe.pth'
elif num_layers == 152:
self.model_path = 'data/pretrained_model/resnet152_caffe.pth'
self.dout_base_model = 1024
self.pretrained = pretrained
self.class_agnostic = class_agnostic
self.num_layer = num_layers
self.fix_cnn_base = fix_cnn_base
self.classes = classes
self.n_classes = len(classes)
# loss
self.RCNN_loss_cls = 0
self.RCNN_loss_bbox = 0
#define base
if self.num_layer == 101:
resnet = resnet101()
elif self.num_layer == 50:
resnet = resnet50()
elif self.num_layer == 152:
resnet = resnet152()
if self.pretrained:
print("Step1: Loading pretrained weights from %s" %(self.model_path))
state_dict = torch.load(self.model_path)
resnet.load_state_dict({k:v for k,v in state_dict.items() if k in resnet.state_dict()})
# Build resnet.
self.RCNN_base = nn.Sequential(resnet.conv1, resnet.bn1,resnet.relu, \
resnet.maxpool,resnet.layer1,resnet.layer2,resnet.layer3)
# Fix blocks
if fix_cnn_base:
for p in self.RCNN_base[0].parameters(): p.requires_grad=False
for p in self.RCNN_base[1].parameters(): p.requires_grad=False
assert (0 <= cfg.RESNET.FIXED_BLOCKS < 4)
if cfg.RESNET.FIXED_BLOCKS >= 3:
for p in self.RCNN_base[6].parameters(): p.requires_grad=False
if cfg.RESNET.FIXED_BLOCKS >= 2:
for p in self.RCNN_base[5].parameters(): p.requires_grad=False
if cfg.RESNET.FIXED_BLOCKS >= 1:
for p in self.RCNN_base[4].parameters(): p.requires_grad=False
def set_bn_fix(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
for p in m.parameters(): p.requires_grad=False
self.RCNN_base.apply(set_bn_fix)
# define rpn
self.RCNN_rpn = _RPN(self.dout_base_model)
#init weight
def normal_init(m, mean, stddev, truncated=False):
"""
weight initalizer: truncated normal and random normal.
"""
# x is a parameter
if truncated:
m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean) # not a perfect approximation
else:
m.weight.data.normal_(mean, stddev)
m.bias.data.zero_()
normal_init(self.RCNN_rpn.RPN_Conv, 0, 0.01, cfg.TRAIN.TRUNCATED)
normal_init(self.RCNN_rpn.RPN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED)
normal_init(self.RCNN_rpn.RPN_bbox_pred, 0, 0.01, cfg.TRAIN.TRUNCATED)
def __init__(self,
classes,
class_agnostic,
pretrained=False,
base_model='vgg16'):
super(_detector, self).__init__()
self.classes = classes
self.n_classes = len(classes)
self.class_agnostic = class_agnostic
self.pretrained = pretrained
#processing of roi get from the rpn
self.RCNN_proposal_target = _ProposalTargetLayer(self.n_classes)
# roi pool, roi align or roicrop
# 1.0/16.0 is because size of input img is 16 times larger than feature map
self.RCNN_roi_pool = _RoIPooling(cfg.POOLING_SIZE, cfg.POOLING_SIZE,
1.0 / 16.0)
self.RCNN_roi_align = RoIAlignAvg(cfg.POOLING_SIZE, cfg.POOLING_SIZE,
1.0 / 16.0)
self.grid_size = cfg.POOLING_SIZE * 2 if cfg.CROP_RESIZE_WITH_MAX_POOL else cfg.POOLING_SIZE
self.RCNN_roi_crop = _RoICrop()
self.base_model = base_model
# top
if base_model == 'vgg16':
self.model_path = 'data/pretrained_model/vgg16_caffe.pth'
vgg = models.vgg16()
if self.pretrained:
state_dict = torch.load(self.model_path)
vgg.load_state_dict({
k: v
for k, v in state_dict.items() if k in vgg.state_dict()
})
self.RCNN_top = nn.Sequential(
*list(vgg.classifier._modules.values())[:-1])
# not using the last maxpool layer
self.RCNN_cls_score = nn.Linear(4096, self.n_classes)
if self.class_agnostic:
self.RCNN_bbox_pred = nn.Linear(4096, 4)
else:
self.RCNN_bbox_pred = nn.Linear(4096, 4 * self.n_classes)
elif base_model == 'resnet101' or base_model == 'resnet50' or base_model == 'resnet152':
if base_model == 'resnet101':
self.model_path = 'data/pretrained_model/resnet101_caffe.pth'
resnet = resnet101()
elif base_model == 'resnet50':
self.model_path = 'data/pretrained_model/resnet50_caffe.pth'
resnet = resnet50()
elif base_model == 'resnet152':
self.model_path = 'data/pretrained_model/resnet152_caffe.pth'
resnet = resnet152()
if self.pretrained:
print("Detector: Loading pretrained weights from %s" %
(self.model_path))
state_dict = torch.load(self.model_path)
resnet.load_state_dict({
k: v
for k, v in state_dict.items() if k in resnet.state_dict()
})
self.RCNN_top = nn.Sequential(resnet.layer4)
self.RCNN_cls_score = nn.Linear(2048, self.n_classes)
if self.class_agnostic:
self.RCNN_bbox_pred = nn.Linear(2048, 4)
else:
self.RCNN_bbox_pred = nn.Linear(2048, 4 * self.n_classes)
def set_bn_fix(m):
classname = m.__class__.__name__
if classname.find('BatchNorm') != -1:
for p in m.parameters():
p.requires_grad = False
self.RCNN_top.apply(set_bn_fix)
else:
print("no support for other CNN model")
exit()