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
