def eval_net():
# Assign imdb_name and imdbval_name according to args.dataset.
if args.dataset == "voc":
args.imdb_name = "voc_2007_trainval+voc_2012_trainval"
args.imdbval_name = "voc_2007_test"
elif args.dataset == "coco":
args.imdb_name = "coco_2014_train+coco_2014_valminusminival"
args.imdbval_name = "coco_2014_minival"
# Import config
if args.dataset == 'coco':
cfg = (coco320, coco512)[args.input_size == 512]
elif args.dataset == 'voc':
cfg = (voc320, voc512)[args.input_size == 512]
# Create imdb, roidb and blob_dataset
print('Create or load an evaluted imdb.')
imdb, roidb = combined_roidb(args.imdbval_name, False)
imdb.competition_mode(on=True)
print('{:d} roidb entries'.format(len(roidb)))
blob_dataset = BlobDataset(imdb,
roidb,
transform=BaseTransform(cfg['min_dim'], MEANS),
target_normalization=True)
# Construct networks.
print('Construct {}_refinedet network.'.format(args.network))
if args.network == 'vgg16':
refinedet = VGGRefineDet(cfg['num_classes'], cfg)
elif args.network == 'resnet101':
refinedet = ResNetRefineDet(cfg['num_classes'], cfg)
refinedet.create_architecture()
# For CPU
net = refinedet
# For GPU/GPUs
if args.cuda:
if num_gpus > 1:
net = torch.nn.DataParallel(refinedet)
else:
net = refinedet.cuda()
cudnn.benchmark = True
# Load weights
net.load_weights(args.model_path)
net.eval()
print('Test RefineDet on:', args.imdbval_name)
print('Using the specified args:')
print(args)
#data_loader = data.DataLoader(blob_dataset,
# batch_size=1,
# num_workers=0,
# shuffle=False,
# collate_fn=detection_collate,
# pin_memory=True)
# all detections are collected into:
# all_boxes[cls][image] = N x 5 array of detections in
# (x1, y1, x2, y2, score)
num_images = len(imdb.image_index)
# num_images = len(blob_dataset)
num_classes = imdb.num_classes
# num_object_classes + 1 ?
print(num_classes)
all_boxes = [[[] for _ in range(num_images)] for _ in range(num_classes)]
empty_array = np.transpose(np.array([[], [], [], [], []]), (1, 0))
output_dir = get_output_dir(imdb, args.result_path)
_t = {'im_detect': Timer(), 'misc': Timer()}
det_file = os.path.join(output_dir, 'detections.pkl')
# pdb.set_trace()
for idx in range(num_images):
img, gt, h, w = blob_dataset.pull_item(idx)
input = Variable(img.unsqueeze(0), volatile=True)
if args.cuda:
input = input.cuda()
# timers forward
_t['im_detect'].tic()
#start = time.time()
# pdb.set_trace()
detection = net(input)
detect_time = _t['im_detect'].toc(average=True)
#detect_time = time.time() - start
print('im_detect: {:d}/{:d} {:.3f}s\n'.format(idx + 1, num_images,
detect_time))
# skip jc = 0, because it's the background class
for jc in range(1, num_classes):
dets = detection[0, jc, :]
mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t()
dets = torch.masked_select(dets, mask).view(-1, 5)
if dets.dim() > 0:
boxes = dets[:, 1:]
boxes[:, 0] *= w
boxes[:, 2] *= w
boxes[:, 1] *= h
boxes[:, 3] *= h
scores = dets[:, 0].cpu().numpy()
cls_dets = np.hstack((boxes.cpu().numpy(),
scores[:,
np.newaxis])).astype(np.float32,
copy=False)
all_boxes[jc][idx] = cls_dets
else:
all_boxes[jc][idx] = empty_array
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
print('Evaluating detections')
imdb.evaluate_detections(all_boxes, output_dir)
def train():
# Assign imdb_name and imdbval_name according to args.dataset.
if args.dataset == "voc":
#args.imdb_name = "voc_2007_trainval"
args.imdb_name = "voc_2007_trainval+voc_2012_trainval"
args.imdbval_name = "voc_2007_test"
elif args.dataset == "coco":
args.imdb_name = "coco_2014_train+coco_2014_valminusminival"
args.imdbval_name = "coco_2014_minival"
# Import config
if args.dataset == 'coco':
cfg = (coco320, coco512)[args.input_size == 512]
elif args.dataset == 'voc':
cfg = (voc320, voc512)[args.input_size == 512]
# Create imdb, roidb and blob_dataset
print('Create or load an imdb.')
imdb, roidb = combined_roidb(args.imdb_name)
blob_dataset = BlobDataset(imdb,
roidb,
transform=SSDAugmentation(
cfg['min_dim'], MEANS),
target_normalization=True)
# Construct networks.
print('Construct {}_refinedet network.'.format(args.network))
if args.network == 'vgg16':
refinedet = VGGRefineDet(cfg['num_classes'], cfg)
elif args.network == 'resnet101':
refinedet = ResNetRefineDet(cfg['num_classes'], cfg)
refinedet.create_architecture(os.path.join(args.save_folder, args.basenet),
pretrained=True,
fine_tuning=True)
#pdb.set_trace()
# For CPU
net = refinedet
# For GPU/GPUs
if args.cuda:
if num_gpus > 1:
net = torch.nn.DataParallel(refinedet)
else:
net = refinedet.cuda()
cudnn.benchmark = True
# Resume
if args.resume_checkpoint:
print('Resuming training, loading {}...'.format(
args.resume_checkpoint))
net.load_weights(args.resume_checkpoint)
# pdb.set_trace()
# params = net.state_dict()
# for k, v in params.items():
# print(k)
# print(v.shape)
optimizer = optim.SGD(filter(lambda p: p.requires_grad, net.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
net.train()
print('Training RefineDet on:', args.imdb_name)
print('Using the specified args:')
print(args)
step_index = 0
str_input_size = str(cfg['min_dim'])
model_info = 'refinedet{0}_'.format(str_input_size) + args.dataset
model_save_folder = os.path.join(args.save_folder, model_info)
if not os.path.exists(model_save_folder):
os.mkdir(model_save_folder)
data_loader = data.DataLoader(blob_dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
# Create batch iterator
# Number of iterations in each epoch
num_iter_per_epoch = len(blob_dataset) // args.batch_size
# number of epoch, in case of resuming from args.start_iter
num_epoch = (cfg['max_iter'] - args.start_iter) // num_iter_per_epoch
iteration = args.start_iter
arm_loss_loc = 0
arm_loss_conf = 0
odm_loss_loc = 0
odm_loss_conf = 0
for epoch in range(0, num_epoch):
# pdb.set_trace()
for i_batch, (images, targets) in enumerate(data_loader):
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
if args.cuda:
images = Variable(images.cuda())
targets = Variable(targets.cuda())
else:
images = Variable(images)
targets = Variable(targets)
# forward
t0 = time.time()
# backprop
optimizer.zero_grad()
bi_loss_loc, bi_loss_conf, multi_loss_loc, multi_loss_conf = \
net(images, targets)
loss = bi_loss_loc.mean() + bi_loss_conf.mean() + \
multi_loss_loc.mean() + multi_loss_conf.mean()
loss.backward()
optimizer.step()
t1 = time.time()
if num_gpus > 1:
arm_loss_loc = bi_loss_loc.mean().data[0]
arm_loss_conf = bi_loss_conf.mean().data[0]
odm_loss_loc = multi_loss_loc.mean().data[0]
odm_loss_conf = multi_loss_conf.mean().data[0]
else:
arm_loss_loc = bi_loss_loc.data[0]
arm_loss_conf = bi_loss_conf.data[0]
odm_loss_loc = multi_loss_loc.data[0]
odm_loss_conf = multi_loss_conf.data[0]
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) +
(' || ARM Loss Loc: %.4f || ARM Loss Conf: %.4f' +
' || ODM Loss Loc: %.4f || ODM Loss Conf: %.4f' +
' || Loss: %.4f ||') %
(arm_loss_loc, arm_loss_conf, odm_loss_loc,
odm_loss_conf, loss.data[0]) + ' ')
# print('iter ' + repr(iteration) +
# ' || Loss: %.4f ||' % (loss.data[0]) + ' ')
# print('iter ' + repr(iteration) +
# ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
if iteration != 0 and iteration % 10000 == 0:
#if iteration != 0 and iteration % cfg['checkpoint_step'] == 0:
print('Saving state, iter:', iteration)
torch.save(
refinedet.state_dict(),
os.path.join(model_save_folder,
model_info + '_' + repr(iteration) + '.pth'))
iteration += 1
torch.save(refinedet.state_dict(),
os.path.join(args.save_folder, args.dataset + '.pth'))