num_boxes = Variable(num_boxes)
gt_boxes = Variable(gt_boxes)
if args.cuda:
cfg.CUDA = True
if args.lighthead:
lighthead = True
# initilize the network here.
if args.net == 'res101':
_RCNN = resnet(imdb.classes, 101, pretrained=True, class_agnostic=args.class_agnostic, lighthead=lighthead)
elif args.net == 'xception':
_RCNN = xception(imdb.classes, pretrained=False, class_agnostic=args.class_agnostic, lighthead=lighthead)
elif args.net == 'mobilenet':
_RCNN = mobilenetv2(imdb.classes, pretrained=True, class_agnostic=args.class_agnostic, lighthead=lighthead)
_RCNN.create_architecture()
lr = cfg.TRAIN.LEARNING_RATE
lr = args.lr
# tr_momentum = cfg.TRAIN.MOMENTUM
# tr_momentum = args.momentum
params = []
for key, value in dict(_RCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
params += [{'params': [value], 'lr': lr * (cfg.TRAIN.DOUBLE_BIAS + 1), \
'weight_decay': cfg.TRAIN.BIAS_DECAY and cfg.TRAIN.WEIGHT_DECAY or 0}]
else:
lighthead=lighthead)
elif args.net == 'squeeze1_0':
_RCNN = squeezenet(pascal_classes,
version='1_0',
pretrained=False,
class_agnostic=args.class_agnostic,
lighthead=lighthead)
elif args.net == 'squeeze1_1':
_RCNN = squeezenet(pascal_classes,
version='1_1',
pretrained=False,
class_agnostic=args.class_agnostic,
lighthead=lighthead)
elif args.net == 'mobilenet':
_RCNN = mobilenetv2(pascal_classes,
pretrained=False,
class_agnostic=args.class_agnostic,
lighthead=lighthead)
else:
print("network is not defined")
pdb.set_trace()
_RCNN.create_architecture()
print("load checkpoint %s" % (load_name))
if args.cuda > 0:
checkpoint = torch.load(load_name)
else:
checkpoint = torch.load(load_name,
map_location=(lambda storage, loc: storage))
_RCNN.load_state_dict(checkpoint['model'])