net.resnet.load_state_dict({ k: v for k, v in pretrained_state.items() if k in net.resnet.state_dict() }) for p in net.resnet.conv1.parameters(): p.requires_grad = False for p in net.resnet.bn1.parameters(): p.requires_grad = False for p in net.resnet.layer1.parameters(): p.requires_grad = False for p in net.resnet.layer2.parameters(): p.requires_grad = False params = [] #params = list(net.parameters()) for p in list(net.parameters()): if p.requires_grad == False: continue params.append(p) # optimizer = torch.optim.Adam(params[-8:], lr=lr) optimizer = torch.optim.SGD(params, lr=lr, momentum=momentum, weight_decay=weight_decay) if not os.path.exists(output_dir): os.makedirs(output_dir) # training train_loss = 0 tp, tf, fg, bg = 0., 0., 0, 0 step_cnt = 0
log_interval = cfg.TRAIN.LOG_IMAGE_ITERS # load data imdb = VisualGenome(split=0, num_im=50) roidb = imdb.roidb data_layer = RoIDataLayer(roidb, imdb.num_classes) # load net net = FasterRCNN(classes=imdb.classes, debug=_DEBUG) network.weights_normal_init(net, dev=0.01) network.load_net(pretrained_model, net) # network.load_pretrained_npy(net, 'checkpoints/VGG_imagenet.npy') net.cuda() net.train() params = list(net.parameters()) print("Params are {}".format( '\n'.join(['{}: {}'.format(n, p.size()) for n,p in net.named_parameters()])) ) # optimizer = torch.optim.Adam(params, lr=0.001, eps=1e-4, weight_decay=weight_decay) optimizer = torch.optim.SGD(params[8:], lr=lr, momentum=momentum, weight_decay=weight_decay) if not os.path.exists(output_dir): os.mkdir(output_dir) # tensorboad use_tensorboard = use_tensorboard and CrayonClient is not None if use_tensorboard: cc = CrayonClient(hostname='127.0.0.1') if remove_all_log:
def main(): global args print "Loading training set and testing set..." train_set = visual_genome(args.dataset_option, 'train') test_set = visual_genome(args.dataset_option, 'test') object_classes = test_set.object_classes print "Done." train_loader = torch.utils.data.DataLoader(train_set, batch_size=1, shuffle=True, num_workers=8, pin_memory=True) test_loader = torch.utils.data.DataLoader(test_set, batch_size=1, shuffle=False, num_workers=8, pin_memory=True) net = FasterRCNN(nhidden=args.mps_feature_len, use_kmeans_anchors=args.use_kmeans_anchors, n_classes=len(object_classes), model=args.base_model) if args.resume_model: print 'Resume training from: {}'.format(args.resume_model) if len(args.resume_model) == 0: raise Exception('[resume_model] not specified') network.load_net(args.detection_model, net) # optimizer = torch.optim.SGD([ # {'params': list(net.parameters())}, # ], lr=args.lr, momentum=args.momentum, weight_decay=0.0005) else: print 'Training from scratch...Initializing network...' optimizer = torch.optim.SGD(list(net.parameters()), lr=args.lr, momentum=args.momentum, weight_decay=0.0005) # network.set_trainable(net.features, requires_grad=True) net.cuda() if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) best_map = 0.0 for epoch in range(0, args.max_epoch): # Training train(train_loader, net, optimizer, epoch) # update learning rate if epoch % args.step_size == args.step_size - 1: args.clip_gradient = False args.lr /= 5 for param_group in optimizer.param_groups: param_group['lr'] = args.lr save_name = os.path.join( args.output_dir, '{}_epoch_{}.h5'.format(args.model_name, epoch)) network.save_net(save_name, net) print('save model: {}'.format(save_name)) try: # Testing map = evaluate(test_loader, net, object_classes) print( 'Epoch[{epoch:d}]: ' 'Recall: ' 'object: {map: .3f}%% (Best: {best_map: .3f}%%)'.format( epoch=epoch, map=map * 100, best_map=best_map * 100)) if map > best_map: best_map = map save_name = os.path.join( args.output_dir, '{}_best.h5'.format(args.model_name, epoch)) network.save_net(save_name, net) except: continue