def main():
args = _parse_args()
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.opts is not None:
cfg_from_list(args.opts)
assert_and_infer_cfg()
test_output_dir = get_output_dir(training=False)
json_data, _, _, _, _ = get_roidb_and_dataset(None, include_gt=True)
run_posetrack_tracking(test_output_dir, json_data)
def main():
args = _parse_args()
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.opts is not None:
cfg_from_list(args.opts)
assert_and_infer_cfg()
test_output_dir = get_output_dir(training=False)
roidb, dataset, _, _, _ = get_roidb_and_dataset(None)
run_mpii_eval(test_output_dir, roidb, dataset)
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True, device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
for i in xrange(num_images):
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2
)
if args.merge_pdfs and num_images > 1:
merge_out_path = '{}/results.pdf'.format(args.output_dir)
if os.path.exists(merge_out_path):
os.remove(merge_out_path)
command = "pdfunite {}/*.pdf {}".format(args.output_dir,
merge_out_path)
subprocess.call(command, shell=True)
parser.print_help()
sys.exit(1)
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
print('Called with args:')
print(args)
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
if args.exp_dir is not None:
cfg.EXP_DIR = args.exp_dir
cfg.GPU_ID = args.gpu_id
print('Using config:')
pprint.pprint(cfg)
caffe.set_mode_gpu()
caffe.set_device(args.gpu_id)
output_dir_name = 'test'
if args.datasets:
output_dir_name += '_' + '_'.join(args.datasets)
output_dir_name += '_' + datetime.datetime.now().strftime("%d_%m_%Y_%H_%M")
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train',)
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train',)
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH
original_num_gpus = cfg.NUM_GPUS
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
effective_batch_size = args.iter_size * args.batch_size
print('effective_batch_size = batch_size * iter_size = %d * %d' % (args.batch_size, args.iter_size))
print('Adaptive config changes:')
print(' effective_batch_size: %d --> %d' % (original_batch_size, effective_batch_size))
print(' NUM_GPUS: %d --> %d' % (original_num_gpus, cfg.NUM_GPUS))
print(' IMS_PER_BATCH: %d --> %d' % (original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH))
### Adjust learning based on batch size change linearly
# For iter_size > 1, gradients are `accumulated`, so lr is scaled based
# on batch_size instead of effective_batch_size
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch_size change:\n'
' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR))
### Adjust solver steps
step_scale = original_batch_size / effective_batch_size
old_solver_steps = cfg.SOLVER.STEPS
old_max_iter = cfg.SOLVER.MAX_ITER
cfg.SOLVER.STEPS = list(map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS))
cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5)
print('Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n'
' SOLVER.STEPS: {} --> {}\n'
' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps, cfg.SOLVER.STEPS,
old_max_iter, cfg.SOLVER.MAX_ITER))
# Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function
# of `collect_and_distribute_fpn_rpn_proposals.py`
#
# post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5)
if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN:
cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch
print('Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n'
' cfg.FPN.RPN_COLLECT_SCALE: {}'.format(cfg.FPN.RPN_COLLECT_SCALE))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
assert_and_infer_cfg()
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time)
# Effective training sample size for one epoch
train_size = roidb_size // args.batch_size * args.batch_size
batchSampler = BatchSampler(
sampler=MinibatchSampler(ratio_list, ratio_index),
batch_size=args.batch_size,
drop_last=True
)
dataset = RoiDataLoader(
roidb,
cfg.MODEL.NUM_CLASSES,
training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_sampler=batchSampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
dataiterator = iter(dataloader)
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
gn_param_nameset = set()
for name, module in maskRCNN.named_modules():
if isinstance(module, nn.GroupNorm):
gn_param_nameset.add(name+'.weight')
gn_param_nameset.add(name+'.bias')
gn_params = []
gn_param_names = []
bias_params = []
bias_param_names = []
nonbias_params = []
nonbias_param_names = []
nograd_param_names = []
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
bias_param_names.append(key)
elif key in gn_param_nameset:
gn_params.append(value)
gn_param_names.append(key)
else:
nonbias_params.append(value)
nonbias_param_names.append(key)
else:
nograd_param_names.append(key)
assert (gn_param_nameset - set(nograd_param_names) - set(bias_param_names)) == set(gn_param_names)
# Learning rate of 0 is a dummy value to be set properly at the start of training
params = [
{'params': nonbias_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY},
{'params': bias_params,
'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0},
{'params': gn_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY_GN}
]
# names of paramerters for each paramter
param_names = [nonbias_param_names, bias_param_names, gn_param_names]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
args.start_step = checkpoint['step'] + 1
if 'train_size' in checkpoint: # For backward compatibility
if checkpoint['train_size'] != train_size:
print('train_size value: %d different from the one in checkpoint: %d'
% (train_size, checkpoint['train_size']))
# reorder the params in optimizer checkpoint's params_groups if needed
# misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint)
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer'])
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0]['lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name() + '_step'
output_dir = misc_utils.get_output_dir(args, args.run_name)
args.cfg_filename = os.path.basename(args.cfg_file)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS)
# Set index for decay steps
decay_steps_ind = None
for i in range(1, len(cfg.SOLVER.STEPS)):
if cfg.SOLVER.STEPS[i] >= args.start_step:
decay_steps_ind = i
break
if decay_steps_ind is None:
decay_steps_ind = len(cfg.SOLVER.STEPS)
training_stats = TrainingStats(
args,
args.disp_interval,
tblogger if args.use_tfboard and not args.no_save else None)
try:
logger.info('Training starts !')
step = args.start_step
for step in range(args.start_step, cfg.SOLVER.MAX_ITER):
# Warm up
if step < cfg.SOLVER.WARM_UP_ITERS:
method = cfg.SOLVER.WARM_UP_METHOD
if method == 'constant':
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR
elif method == 'linear':
alpha = step / cfg.SOLVER.WARM_UP_ITERS
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 - alpha) + alpha
else:
raise KeyError('Unknown SOLVER.WARM_UP_METHOD: {}'.format(method))
lr_new = cfg.SOLVER.BASE_LR * warmup_factor
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
elif step == cfg.SOLVER.WARM_UP_ITERS:
net_utils.update_learning_rate(optimizer, lr, cfg.SOLVER.BASE_LR)
lr = optimizer.param_groups[0]['lr']
assert lr == cfg.SOLVER.BASE_LR
# Learning rate decay
if decay_steps_ind < len(cfg.SOLVER.STEPS) and \
step == cfg.SOLVER.STEPS[decay_steps_ind]:
logger.info('Decay the learning on step %d', step)
lr_new = lr * cfg.SOLVER.GAMMA
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
decay_steps_ind += 1
training_stats.IterTic()
optimizer.zero_grad()
for inner_iter in range(args.iter_size):
try:
input_data = next(dataiterator)
except StopIteration:
dataiterator = iter(dataloader)
input_data = next(dataiterator)
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
net_outputs = maskRCNN(**input_data)
training_stats.UpdateIterStats(net_outputs, inner_iter)
loss = net_outputs['total_loss']
loss.backward()
optimizer.step()
training_stats.IterToc()
training_stats.LogIterStats(step, lr)
if (step+1) % CHECKPOINT_PERIOD == 0:
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
# ---- Training ends ----
# Save last checkpoint
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
except (RuntimeError, KeyboardInterrupt):
del dataiterator
logger.info('Save ckpt on exception ...')
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
logger.info('Save ckpt done.')
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
if args.use_tfboard and not args.no_save:
tblogger.close()
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.img_dir or args.images
assert bool(args.img_dir) ^ bool(args.images)
prefix_path = args.output_dir
os.makedirs(prefix_path, exist_ok=True)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
# with open(args.img_list, 'rb') as f:
# imglist = pickle.load(f)
imglist = file_name(args.img_dir)
num_images = len(imglist)
print('num_images: ', num_images)
writen_results = []
# # validate
# demo_im = cv2.imread(imglist[0])
# print(np.shape(demo_im))
# h, w, _ = np.shape(demo_im)
# #print(h)
# #print(args.height)
# assert h == args.height
# assert w == args.width
# h_scale = 720 / args.height
# w_scale = 1280 / args.width
for i in tqdm(range(num_images)):
im = cv2.imread(
'/home/xinyu/dataset/Exclusively-Dark-Image-Dataset/ExDark/' +
imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
# boxs = [[x1, y1, x2, y2, cls], ...]
boxes, _, _, classes = convert_from_cls_format(cls_boxes, cls_segms,
cls_keyps)
if boxes is None:
continue
# scale
boxes[:, 0] = boxes[:, 0] #* w_scale
boxes[:, 2] = boxes[:, 2] #* w_scale
boxes[:, 1] = boxes[:, 1] #* h_scale
boxes[:, 3] = boxes[:, 3] #* h_scale
if classes == []:
continue
for instance_idx, cls_idx in enumerate(classes):
cls_name = dataset.classes[cls_idx]
if cls_name == 'bicycle':
cls_name = 'Bicycle'
elif cls_name == 'dog':
cls_name = 'Dog'
elif cls_name == 'boat':
cls_name = 'Boat'
elif cls_name == 'bottle':
cls_name = 'Bottle'
elif cls_name == 'bus':
cls_name = 'Bus'
elif cls_name == 'car':
cls_name = 'Car'
elif cls_name == 'cat':
cls_name = 'Cat'
elif cls_name == 'chair':
cls_name = 'Chair'
elif cls_name == 'cup':
cls_name = 'Cup'
elif cls_name == 'motorcycle':
cls_name = 'Motorbike'
elif cls_name == 'person':
cls_name = 'People'
elif cls_name == 'dining table':
cls_name = 'Table'
if cls_name not in edark_category:
continue
writen_results.append({
"name": imglist[i].split('/'),
"timestamp": 1000,
"category": cls_name,
"bbox": boxes[instance_idx, :4],
"score": boxes[instance_idx, -1]
})
with open(os.path.join(prefix_path, args.name + '.json'),
'w') as outputfile:
json.dump(writen_results, outputfile, cls=MyEncoder)
combined = np.hstack((gt, pred))
im_name = entry['image']
if isinstance(im_name, list):
im_name = im_name[len(im_name) // 2]
out_name = im_name[len(dataset.image_directory):]
out_path = osp.join(output_dir, out_name)
gen_utils.mkdir_p(osp.dirname(out_path))
cv2.imwrite(out_path, combined)
if __name__ == '__main__':
args = _parse_args()
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.opts is not None:
cfg_from_list(args.opts)
assert_and_infer_cfg()
test_output_dir = get_output_dir(training=False)
det_file = osp.join(test_output_dir, 'detections.pkl')
tracking_det_file = osp.join(test_output_dir, 'detections_withTracks.pkl')
if osp.exists(tracking_det_file):
det_file = tracking_det_file
output_dir = osp.join(test_output_dir, 'vis/')
if not osp.exists(det_file):
raise ValueError('Output file not found {}'.format(det_file))
else:
logger.info('Visualizing {}'.format(det_file))
# Set include_gt True when using the roidb to evalute directly. Not doing
# that currently
roidb, dataset, _, _, _ = get_roidb_and_dataset(None, include_gt=True)
vis(roidb, det_file, args.thresh, output_dir)
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "keypoints_carfusion_butler":
cfg.TRAIN.DATASETS = ('keypoints_carfusion_butler', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "keypoints_carfusion_fifth":
cfg.TRAIN.DATASETS = ('keypoints_carfusion_fifth', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "keypoints_carfusion_craig":
cfg.TRAIN.DATASETS = ('keypoints_carfusion_craig', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "keypoints_carfusion_morewood":
cfg.TRAIN.DATASETS = ('keypoints_carfusion_morewood', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "keypoints_carfusion":
cfg.TRAIN.DATASETS = ('keypoints_carfusion', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "keypoints_carfusion_coco":
cfg.TRAIN.DATASETS = ('keypoints_carfusion_coco', )
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH
original_num_gpus = cfg.NUM_GPUS
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
effective_batch_size = args.iter_size * args.batch_size
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(args.batch_size, args.iter_size))
print('Adaptive config changes:')
print(' effective_batch_size: %d --> %d' %
(original_batch_size, effective_batch_size))
print(' NUM_GPUS: %d --> %d' %
(original_num_gpus, cfg.NUM_GPUS))
print(' IMS_PER_BATCH: %d --> %d' %
(original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH))
### Adjust learning based on batch size change linearly
# For iter_size > 1, gradients are `accumulated`, so lr is scaled based
# on batch_size instead of effective_batch_size
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch_size change:\n'
' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR))
### Adjust solver steps
step_scale = original_batch_size / effective_batch_size
old_solver_steps = cfg.SOLVER.STEPS
old_max_iter = cfg.SOLVER.MAX_ITER
cfg.SOLVER.STEPS = list(
map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS))
cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5)
print(
'Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n'
' SOLVER.STEPS: {} --> {}\n'
' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps,
cfg.SOLVER.STEPS, old_max_iter,
cfg.SOLVER.MAX_ITER))
# Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function
# of `collect_and_distribute_fpn_rpn_proposals.py`
#
# post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5)
if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN:
cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch
print(
'Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n'
' cfg.FPN.RPN_COLLECT_SCALE: {}'.format(
cfg.FPN.RPN_COLLECT_SCALE))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
assert_and_infer_cfg()
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb',
timers['roidb'].average_time)
# Effective training sample size for one epoch
train_size = roidb_size // args.batch_size * args.batch_size
sampler = MinibatchSampler(ratio_list, ratio_index)
dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
drop_last=True,
sampler=sampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
dataiterator = iter(dataloader)
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
bias_params = []
bias_param_names = []
nonbias_params = []
nonbias_param_names = []
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
bias_param_names.append(key)
else:
nonbias_params.append(value)
nonbias_param_names.append(key)
# Learning rate of 0 is a dummy value to be set properly at the start of training
params = [{
'params': nonbias_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY
}, {
'params':
bias_params,
'lr':
0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay':
cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0
}]
# names of paramerters for each paramter
param_names = [nonbias_param_names, bias_param_names]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
args.start_step = checkpoint['step'] + 1
if 'train_size' in checkpoint: # For backward compatibility
if checkpoint['train_size'] != train_size:
print(
'train_size value: %d different from the one in checkpoint: %d'
% (train_size, checkpoint['train_size']))
# reorder the params in optimizer checkpoint's params_groups if needed
# misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint)
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer,
checkpoint['optimizer'])
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0][
'lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name() + '_step'
output_dir = misc_utils.get_output_dir(args, args.run_name)
args.cfg_filename = os.path.basename(args.cfg_file)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS)
# Set index for decay steps
decay_steps_ind = None
for i in range(1, len(cfg.SOLVER.STEPS)):
if cfg.SOLVER.STEPS[i] >= args.start_step:
decay_steps_ind = i
break
if decay_steps_ind is None:
decay_steps_ind = len(cfg.SOLVER.STEPS)
training_stats = TrainingStats(
args, args.disp_interval,
tblogger if args.use_tfboard and not args.no_save else None)
try:
logger.info('Training starts !')
step = args.start_step
for step in range(args.start_step, cfg.SOLVER.MAX_ITER):
# Warm up
if step < cfg.SOLVER.WARM_UP_ITERS:
method = cfg.SOLVER.WARM_UP_METHOD
if method == 'constant':
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR
elif method == 'linear':
alpha = step / cfg.SOLVER.WARM_UP_ITERS
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 -
alpha) + alpha
else:
raise KeyError(
'Unknown SOLVER.WARM_UP_METHOD: {}'.format(method))
lr_new = cfg.SOLVER.BASE_LR * warmup_factor
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
elif step == cfg.SOLVER.WARM_UP_ITERS:
net_utils.update_learning_rate(optimizer, lr,
cfg.SOLVER.BASE_LR)
lr = optimizer.param_groups[0]['lr']
assert lr == cfg.SOLVER.BASE_LR
# Learning rate decay
if decay_steps_ind < len(cfg.SOLVER.STEPS) and \
step == cfg.SOLVER.STEPS[decay_steps_ind]:
logger.info('Decay the learning on step %d', step)
lr_new = lr * cfg.SOLVER.GAMMA
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
decay_steps_ind += 1
training_stats.IterTic()
optimizer.zero_grad()
for inner_iter in range(args.iter_size):
try:
input_data = next(dataiterator)
except StopIteration:
dataiterator = iter(dataloader)
input_data = next(dataiterator)
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
net_outputs = maskRCNN(**input_data)
training_stats.UpdateIterStats(net_outputs, inner_iter)
loss = net_outputs['total_loss']
loss.backward()
optimizer.step()
training_stats.IterToc()
training_stats.LogIterStats(step, lr)
if (step + 1) % CHECKPOINT_PERIOD == 0:
save_ckpt(output_dir, args, step, train_size, maskRCNN,
optimizer)
# ---- Training ends ----
# Save last checkpoint
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
except (RuntimeError, KeyboardInterrupt):
del dataiterator
logger.info('Save ckpt on exception ...')
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
logger.info('Save ckpt done.')
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
if args.use_tfboard and not args.no_save:
tblogger.close()
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "rsna2018":
cfg.TRAIN.DATASETS = ('RSNA_2018_pos_train',)
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train',)
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_val',)
cfg.MODEL.NUM_CLASSES = 81
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
print('Batch size change from {} (in config file) to {}'.format(
original_batch_size, args.batch_size))
print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' % (cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Adjust learning based on batch size change linearly
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch size change: {} --> {}'.format(
old_base_lr, cfg.SOLVER.BASE_LR))
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
train_size = len(roidb)
logger.info('{:d} roidb entries'.format(train_size))
logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time)
sampler = MinibatchSampler(ratio_list, ratio_index)
dataset = RoiDataLoader(
roidb,
cfg.MODEL.NUM_CLASSES,
training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
sampler=sampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
assert_and_infer_cfg()
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
bias_params = []
nonbias_params = []
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
else:
nonbias_params.append(value)
params = [
{'params': nonbias_params,
'lr': cfg.SOLVER.BASE_LR,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY},
{'params': bias_params,
'lr': cfg.SOLVER.BASE_LR * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0}
]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
assert checkpoint['iters_per_epoch'] == train_size // args.batch_size, \
"iters_per_epoch should match for resume"
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer'])
if checkpoint['step'] == (checkpoint['iters_per_epoch'] - 1):
# Resume from end of an epoch
args.start_epoch = checkpoint['epoch'] + 1
args.start_iter = 0
else:
# Resume from the middle of an epoch.
# NOTE: dataloader is not synced with previous state
args.start_epoch = checkpoint['epoch']
args.start_iter = checkpoint['step'] + 1
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0]['lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name()
output_dir = misc_utils.get_output_dir(args, args.run_name)
args.cfg_filename = os.path.basename(args.cfg_file)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
training_stats = TrainingStats(
args,
args.disp_interval,
tblogger if args.use_tfboard and not args.no_save else None)
iters_per_epoch = int(train_size / args.batch_size) # drop last
args.iters_per_epoch = iters_per_epoch
ckpt_interval_per_epoch = iters_per_epoch // args.ckpt_num_per_epoch
try:
logger.info('Training starts !')
args.step = args.start_iter
global_step = iters_per_epoch * args.start_epoch + args.step
for args.epoch in range(args.start_epoch, args.start_epoch + args.num_epochs):
# ---- Start of epoch ----
# adjust learning rate
if args.lr_decay_epochs and args.epoch == args.lr_decay_epochs[0] and args.start_iter == 0:
args.lr_decay_epochs.pop(0)
net_utils.decay_learning_rate(optimizer, lr, cfg.SOLVER.GAMMA)
lr *= cfg.SOLVER.GAMMA
for args.step, input_data in zip(range(args.start_iter, iters_per_epoch), dataloader):
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
training_stats.IterTic()
net_outputs = maskRCNN(**input_data)
training_stats.UpdateIterStats(net_outputs)
loss = net_outputs['total_loss']
optimizer.zero_grad()
loss.backward()
optimizer.step()
training_stats.IterToc()
if (args.step+1) % ckpt_interval_per_epoch == 0:
net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer)
if args.step % args.disp_interval == 0:
log_training_stats(training_stats, global_step, lr)
global_step += 1
# ---- End of epoch ----
# save checkpoint
net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer)
# reset starting iter number after first epoch
args.start_iter = 0
# ---- Training ends ----
if iters_per_epoch % args.disp_interval != 0:
# log last stats at the end
log_training_stats(training_stats, global_step, lr)
except (RuntimeError, KeyboardInterrupt):
logger.info('Save ckpt on exception ...')
net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer)
logger.info('Save ckpt done.')
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
if args.use_tfboard and not args.no_save:
tblogger.close()
pseudo_gt_rgb = ambiguous * image_norm + (
1 - ambiguous) * pseudo_gt_rgb
visual.append(pseudo_gt_rgb)
# ready to assemble
visual_logits = torch.cat(visual, -1)
self._visualise_grid(visual_logits, gt_labels, epoch, scores=cls_out)
if __name__ == "__main__":
args = get_arguments(sys.argv[1:])
# Reading the config
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
print("Config: \n", cfg)
trainer = DecTrainer(args)
torch.manual_seed(0)
timer = Timer()
def time_call(func, msg, *args, **kwargs):
timer.reset_stage()
func(*args, **kwargs)
print(msg + (" {:3.2}m".format(timer.get_stage_elapsed() / 60.)))
for epoch in range(trainer.start_epoch, cfg.TRAIN.NUM_EPOCHS + 1):
print("Epoch >>> ", epoch)
def main():
"""main function"""
#Clearn result floder
if os.path.exists('../reid_test/'):
shutil.rmtree('../reid_test/')
os.makedirs('../reid_test/')
if os.path.exists('../output'):
shutil.rmtree('../output/')
os.makedirs('../output/')
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True, device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir) #get the list of test pic
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir) #make the results file
for i in tqdm(xrange(num_images)):
# print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dataset,
box_alpha=0.9,
show_class=True,
thresh=0.1,
kp_thresh=2
)
if args.merge_pdfs and num_images > 1: #option is true
merge_out_path = '{}/results.pdf'.format(args.output_dir)
if os.path.exists(merge_out_path):
os.remove(merge_out_path)
command = "pdfunite {}/*.pdf {}".format(args.output_dir,
merge_out_path)
subprocess.call(command, shell=True)
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "coco2014":
cfg.TRAIN.DATASETS = ('coco_2014_train', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == 'vcoco_trainval':
cfg.TRAIN.DATASETS = ('vcoco_trainval', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == 'vcoco_train':
cfg.TRAIN.DATASETS = ('vcoco_train', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == 'vcoco_val':
cfg.TRAIN.DATASETS = ('vcoco_val', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == 'keypoints_coco2014':
cfg.TRAIN.DATASETS = ('keypoints_coco_2014_train', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
if args.vcoco_kp_on:
cfg.VCOCO.KEYPOINTS_ON = True
cfg.NETWORK_NAME = args.net_name # network name
print('Network name:', args.net_name)
cfg.MODEL.CONV_BODY = args.conv_body # backbone network name
print('Conv_body name:', args.conv_body)
cfg.TRAIN.FG_THRESH = args.fg_thresh
print('Train fg thresh:', args.fg_thresh)
cfg.RESNETS.FREEZE_AT = args.freeze_at
print('Freeze at: ', args.freeze_at)
cfg.VCOCO.MLP_HEAD_DIM = args.mlp_head_dim
print('MLP head dim: ', args.mlp_head_dim)
cfg.SOLVER.MAX_ITER = args.max_iter
print('MAX iter: ', args.max_iter)
cfg.TRAIN.SNAPSHOT_ITERS = args.snapshot
print('Snapshot Iters: ', args.snapshot)
if args.solver_steps is not None:
cfg.SOLVER.STEPS = args.solver_steps
print('Solver_steps: ', cfg.SOLVER.STEPS)
cfg.VCOCO.TRIPLETS_NUM_PER_IM = args.triplets_num_per_im
print('triplets_num_per_im: ', cfg.VCOCO.TRIPLETS_NUM_PER_IM)
cfg.VCOCO.HEATMAP_KERNEL_SIZE = args.heatmap_kernel_size
print('heatmap_kernel_size: ', cfg.VCOCO.HEATMAP_KERNEL_SIZE)
cfg.VCOCO.PART_CROP_SIZE = args.part_crop_size
print('part_crop_size: ', cfg.VCOCO.PART_CROP_SIZE)
print('use use_kps17 for part Align: ', args.use_kps17)
if args.use_kps17:
cfg.VCOCO.USE_KPS17 = True
else:
cfg.VCOCO.USE_KPS17 = False
print('MULTILEVEL_ROIS: ', cfg.FPN.MULTILEVEL_ROIS)
if args.vcoco_use_spatial:
cfg.VCOCO.USE_SPATIAL = True
if args.vcoco_use_union_feat:
cfg.VCOCO.USE_UNION_FEAT = True
if args.use_precomp_box:
cfg.VCOCO.USE_PRECOMP_BOX = True
cfg.DEBUG_TEST_WITH_GT = True
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH # 16
original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH
original_num_gpus = cfg.NUM_GPUS
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
effective_batch_size = args.iter_size * args.batch_size
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(args.batch_size, args.iter_size))
print('Adaptive config changes:')
print(' effective_batch_size: %d --> %d' %
(original_batch_size, effective_batch_size))
print(' NUM_GPUS: %d --> %d' %
(original_num_gpus, cfg.NUM_GPUS))
print(' IMS_PER_BATCH: %d --> %d' %
(original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH))
print(' FG_THRESH: ', cfg.TRAIN.FG_THRESH)
### Adjust learning based on batch size change linearly
# For iter_size > 1, gradients are `accumulated`, so lr is scaled based
# on batch_size instead of effective_batch_size
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch_size change:\n'
' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR))
### Adjust solver steps
step_scale = original_batch_size / effective_batch_size
old_solver_steps = cfg.SOLVER.STEPS
old_max_iter = cfg.SOLVER.MAX_ITER
cfg.SOLVER.STEPS = list(
map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS))
cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5)
cfg.SOLVER.VAL_ITER = int(cfg.SOLVER.VAL_ITER * step_scale + 0.5)
cfg.TRAIN.SNAPSHOT_ITERS = int(cfg.TRAIN.SNAPSHOT_ITERS * step_scale + 0.5)
print(
'Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n'
' SOLVER.STEPS: {} --> {}\n'
' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps,
cfg.SOLVER.STEPS, old_max_iter,
cfg.SOLVER.MAX_ITER))
# Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function
# of `collect_and_distribute_fpn_rpn_proposals.py`
#
# post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5)
if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN:
cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch
print(
'Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n'
' cfg.FPN.RPN_COLLECT_SCALE: {}'.format(
cfg.FPN.RPN_COLLECT_SCALE))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
# ipdb.set_trace()
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
assert_and_infer_cfg()
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb',
timers['roidb'].average_time)
# Effective training sample size for one epoch
train_size = roidb_size // args.batch_size * args.batch_size
# ToDo: shuffle?
batchSampler = BatchSampler(sampler=MinibatchSampler(
ratio_list, ratio_index),
batch_size=args.batch_size,
drop_last=True)
dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_sampler=batchSampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
# dataiterator = iter(dataloader)
### Model ###
from modeling.model_builder import Generalized_RCNN
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
bias_hoi_params = []
bias_hoi_param_names = []
bias_faster_params = []
bias_faster_param_names = []
nobias_hoi_params = []
nobias_hoi_param_names = []
nobias_faster_params = []
nobias_faster_param_names = []
# bias_params = []
# bias_param_names = []
# nonbias_params = []
# nonbias_param_names = []
#base_model = torch.load('Outputs/baseline/baseline_512_32_nogt_1o3/ckpt/model_step47999.pth')
nograd_param_names = []
for key, value in maskRCNN.named_parameters():
#if key in base_model['model'].keys():
# value.requires_grad = False
#print('the key xxx:', key)
# Fix RPN module same as the paper
# ToDo: or key.startswith('Box')
# if 'affinity' not in key:
# value.requires_grad = False
print(key, value.size(), value.requires_grad)
if value.requires_grad:
if 'bias' in key:
if 'HOI_Head' in key:
bias_hoi_params.append(value)
bias_hoi_param_names.append(key)
else:
bias_faster_params.append(value)
bias_faster_param_names.append(key)
else:
if 'HOI_Head' in key:
nobias_hoi_params.append(value)
nobias_hoi_param_names.append(key)
else:
nobias_faster_params.append(value)
nobias_faster_param_names.append(key)
else:
nograd_param_names.append(key)
#del base_model
#ipdb.set_trace()
# Learning rate of 0 is a dummy value to be set properly at the start of training
params = [
{
'params': nobias_hoi_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY
},
{
'params': nobias_faster_params,
'lr': 0 * cfg.SOLVER.FASTER_RCNN_WEIGHT,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY
},
{
'params':
bias_hoi_params,
'lr':
0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay':
cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0
},
{
'params':
bias_faster_params,
'lr':
0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1) *
cfg.SOLVER.FASTER_RCNN_WEIGHT,
'weight_decay':
cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0
},
]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
if args.krcnn_from_faster:
net_utils.load_krcnn_from_faster(maskRCNN, checkpoint['model'])
else:
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
print('Original model loaded....')
if args.resume:
print('Resume, loaded step\n\n\n: ', checkpoint['step'])
args.start_step = checkpoint['step'] + 1
if 'train_size' in checkpoint: # For backward compatibility
if checkpoint['train_size'] != train_size:
print(
'train_size value: %d different from the one in checkpoint: %d'
% (train_size, checkpoint['train_size']))
# reorder the params in optimizer checkpoint's params_groups if needed
# misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint)
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
optimizer.load_state_dict(checkpoint['optimizer'])
# misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer'])
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0][
'lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name() + '_step'
#output_dir = misc_utils.get_output_dir(args, args.run_name)
output_dir = os.path.join('Outputs', args.expDir, args.expID)
os.makedirs(output_dir, exist_ok=True)
args.cfg_filename = os.path.basename(args.cfg_file)
tblogger = None
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
train_val(maskRCNN, args, optimizer, lr, dataloader, train_size,
output_dir, tblogger)
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
cfg.DATASET = args.dataset
if args.dataset == "vg80k":
cfg.TRAIN.DATASETS = ('vg80k_train', )
cfg.TEST.DATASETS = ('vg80k_val', )
cfg.MODEL.NUM_CLASSES = 53305 # includes background
cfg.MODEL.NUM_PRD_CLASSES = 29086 # excludes background
elif args.dataset == "vg8k":
cfg.TRAIN.DATASETS = ('vg8k_train', )
cfg.TEST.DATASETS = ('vg8k_val', )
cfg.MODEL.NUM_CLASSES = 5331 # includes background
cfg.MODEL.NUM_PRD_CLASSES = 2000 # excludes background
elif args.dataset == "vrd":
cfg.TRAIN.DATASETS = ('vrd_train', )
cfg.TEST.DATASETS = ('vrd_val', )
cfg.MODEL.NUM_CLASSES = 101
cfg.MODEL.NUM_PRD_CLASSES = 70 # exclude background
elif args.dataset == "vg":
cfg.TRAIN.DATASETS = ('vg_train', )
cfg.TEST.DATASETS = ('vg_val', )
cfg.MODEL.NUM_CLASSES = 151
cfg.MODEL.NUM_PRD_CLASSES = 50 # exclude background
elif args.dataset == "gvqa20k":
cfg.TRAIN.DATASETS = ('gvqa20k_train', )
cfg.TEST.DATASETS = ('gvqa20k_val', )
cfg.MODEL.NUM_CLASSES = 1704 # includes background
cfg.MODEL.NUM_PRD_CLASSES = 310 # exclude background
elif args.dataset == "gvqa10k":
cfg.TRAIN.DATASETS = ('gvqa10k_train', )
cfg.TEST.DATASETS = ('gvqa10k_val', )
cfg.MODEL.NUM_CLASSES = 1704 # includes background
cfg.MODEL.NUM_PRD_CLASSES = 310 # exclude background
elif args.dataset == "gvqa":
cfg.TRAIN.DATASETS = ('gvqa_train', )
cfg.TEST.DATASETS = ('gvqa_val', )
cfg.MODEL.NUM_CLASSES = 1704 # includes background
cfg.MODEL.NUM_PRD_CLASSES = 310 # exclude background
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
if args.seed:
cfg.RNG_SEED = args.seed
# Some imports need to be done after loading the config to avoid using default values
from datasets.roidb_rel import combined_roidb_for_training
from modeling.model_builder_reltransformer import Generalized_RCNN
from core.test_engine_reltransformer import run_eval_inference, run_inference
from core.test_engine_reltransformer import get_inference_dataset, get_roidb_and_dataset
logger.info('Training with config:')
logger.info(pprint.pformat(cfg))
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH
original_num_gpus = cfg.NUM_GPUS
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
effective_batch_size = args.iter_size * args.batch_size
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(args.batch_size, args.iter_size))
print('Adaptive config changes:')
print(' effective_batch_size: %d --> %d' %
(original_batch_size, effective_batch_size))
print(' NUM_GPUS: %d --> %d' %
(original_num_gpus, cfg.NUM_GPUS))
print(' IMS_PER_BATCH: %d --> %d' %
(original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH))
### Adjust learning based on batch size change linearly
# For iter_size > 1, gradients are `accumulated`, so lr is scaled based
# on batch_size instead of effective_batch_size
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch_size change:\n'
' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR))
### Adjust solver steps
step_scale = original_batch_size / effective_batch_size
old_solver_steps = cfg.SOLVER.STEPS
old_max_iter = cfg.SOLVER.MAX_ITER
cfg.SOLVER.STEPS = list(
map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS))
cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5)
print(
'Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n'
' SOLVER.STEPS: {} --> {}\n'
' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps,
cfg.SOLVER.STEPS, old_max_iter,
cfg.SOLVER.MAX_ITER))
if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN:
cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch
print(
'Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n'
' cfg.FPN.RPN_COLLECT_SCALE: {}'.format(
cfg.FPN.RPN_COLLECT_SCALE))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
assert_and_infer_cfg()
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb',
timers['roidb'].average_time)
# Effective training sample size for one epoch
train_size = roidb_size // args.batch_size * args.batch_size
batchSampler = BatchSampler(sampler=MinibatchSampler(
ratio_list, ratio_index),
batch_size=args.batch_size,
drop_last=True)
dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_sampler=batchSampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
dataiterator = iter(dataloader)
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
# record backbone params, i.e., conv_body and box_head params
gn_params = []
backbone_bias_params = []
backbone_bias_param_names = []
prd_branch_bias_params = []
prd_branch_bias_param_names = []
backbone_nonbias_params = []
backbone_nonbias_param_names = []
prd_branch_nonbias_params = []
prd_branch_nonbias_param_names = []
if cfg.MODEL.DECOUPLE:
for key, value in dict(maskRCNN.named_parameters()).items():
if not 'so_sem_embeddings.2' in key:
value.requires_grad = False
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'gn' in key:
gn_params.append(value)
elif 'Conv_Body' in key or 'Box_Head' in key or 'Box_Outs' in key or 'RPN' in key:
if 'bias' in key:
backbone_bias_params.append(value)
backbone_bias_param_names.append(key)
else:
backbone_nonbias_params.append(value)
backbone_nonbias_param_names.append(key)
else:
if 'bias' in key:
prd_branch_bias_params.append(value)
prd_branch_bias_param_names.append(key)
else:
prd_branch_nonbias_params.append(value)
prd_branch_nonbias_param_names.append(key)
# Learning rate of 0 is a dummy value to be set properly at the start of training
params = [{
'params': backbone_nonbias_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY
}, {
'params':
backbone_bias_params,
'lr':
0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay':
cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0
}, {
'params': prd_branch_nonbias_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY
}, {
'params':
prd_branch_bias_params,
'lr':
0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay':
cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0
}, {
'params': gn_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY_GN
}]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
load_ckpt_dir = './'
### Load checkpoint
if args.load_ckpt_dir:
load_name = get_checkpoint_resume_file(args.load_ckpt_dir)
load_ckpt_dir = args.load_ckpt_dir
elif args.load_ckpt:
load_name = args.load_ckpt
load_ckpt_dir = os.path.dirname(args.load_ckpt)
if args.load_ckpt or args.load_ckpt_dir:
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
if cfg.MODEL.DECOUPLE:
del checkpoint['model']['RelDN.so_sem_embeddings.2.weight']
del checkpoint['model']['RelDN.so_sem_embeddings.2.bias']
del checkpoint['model']['RelDN.prd_sem_embeddings.2.weight']
del checkpoint['model']['RelDN.prd_sem_embeddings.2.bias']
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
args.start_step = checkpoint['step'] + 1
if 'train_size' in checkpoint: # For backward compatibility
if checkpoint['train_size'] != train_size:
print(
'train_size value: %d different from the one in checkpoint: %d'
% (train_size, checkpoint['train_size']))
# reorder the params in optimizer checkpoint's params_groups if needed
# misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint)
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer,
checkpoint['optimizer'])
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[2][
'lr'] # lr of non-backbone parameters, for commmand line outputs.
backbone_lr = optimizer.param_groups[0][
'lr'] # lr of backbone parameters, for commmand line outputs.
prd_categories = maskRCNN.prd_categories
obj_categories = maskRCNN.obj_categories
prd_freq_dict = maskRCNN.prd_freq_dict
obj_freq_dict = maskRCNN.obj_freq_dict
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name() + '_step_with_prd_cls_v' + str(
cfg.MODEL.SUBTYPE)
output_dir = misc_utils.get_output_dir(args, args.run_name)
args.cfg_filename = os.path.basename(args.cfg_file)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
ckpt_dir = os.path.join(output_dir, 'ckpt')
if not os.path.exists(ckpt_dir):
os.makedirs(ckpt_dir)
# if os.path.exists(os.path.join(ckpt_dir, 'best.json')):
# best = json.load(open(os.path.join(ckpt_dir, 'best.json')))
if args.resume and os.path.exists(
os.path.join(load_ckpt_dir, 'best.json')):
logger.info('Loading best json from :' +
os.path.join(load_ckpt_dir, 'best.json'))
best = json.load(open(os.path.join(load_ckpt_dir, 'best.json')))
json.dump(best, open(os.path.join(ckpt_dir, 'best.json'), 'w'))
else:
best = {}
best['avg_per_class_acc'] = 0.0
best['iteration'] = 0
best['accuracies'] = []
json.dump(best, open(os.path.join(ckpt_dir, 'best.json'), 'w'))
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
args.output_dir = output_dir
args.do_val = True
args.use_gt_boxes = True
args.use_gt_labels = True
logger.info('Creating val roidb')
val_dataset_name, val_proposal_file = get_inference_dataset(0)
val_roidb, val_dataset, start_ind, end_ind, total_num_images = get_roidb_and_dataset(
val_dataset_name, val_proposal_file, None, args.do_val)
logger.info('Done')
### Training Loop ###
maskRCNN.train()
# CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS)
# CHECKPOINT_PERIOD = cfg.SOLVER.MAX_ITER / cfg.TRAIN.SNAPSHOT_FREQ
CHECKPOINT_PERIOD = 10000
EVAL_PERIOD = cfg.TRAIN.EVAL_PERIOD
# Set index for decay steps
decay_steps_ind = None
for i in range(1, len(cfg.SOLVER.STEPS)):
if cfg.SOLVER.STEPS[i] >= args.start_step:
decay_steps_ind = i
break
if decay_steps_ind is None:
decay_steps_ind = len(cfg.SOLVER.STEPS)
training_stats = TrainingStats(
args, args.disp_interval,
tblogger if args.use_tfboard and not args.no_save else None)
try:
logger.info('Training starts !')
step = args.start_step
for step in range(args.start_step, cfg.SOLVER.MAX_ITER):
maskRCNN.train()
# Warm up
if step < cfg.SOLVER.WARM_UP_ITERS:
method = cfg.SOLVER.WARM_UP_METHOD
if method == 'constant':
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR
elif method == 'linear':
alpha = step / cfg.SOLVER.WARM_UP_ITERS
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 -
alpha) + alpha
else:
raise KeyError(
'Unknown SOLVER.WARM_UP_METHOD: {}'.format(method))
lr_new = cfg.SOLVER.BASE_LR * warmup_factor
net_utils.update_learning_rate_rel(optimizer, lr, lr_new)
lr = optimizer.param_groups[2]['lr']
backbone_lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
elif step == cfg.SOLVER.WARM_UP_ITERS:
net_utils.update_learning_rate_rel(optimizer, lr,
cfg.SOLVER.BASE_LR)
lr = optimizer.param_groups[2]['lr']
backbone_lr = optimizer.param_groups[0]['lr']
assert lr == cfg.SOLVER.BASE_LR
# Learning rate decay
if decay_steps_ind < len(cfg.SOLVER.STEPS) and \
step == cfg.SOLVER.STEPS[decay_steps_ind]:
logger.info('Decay the learning on step %d', step)
lr_new = lr * cfg.SOLVER.GAMMA
net_utils.update_learning_rate_rel(optimizer, lr, lr_new)
lr = optimizer.param_groups[2]['lr']
backbone_lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
decay_steps_ind += 1
training_stats.IterTic()
optimizer.zero_grad()
for inner_iter in range(args.iter_size):
try:
input_data = next(dataiterator)
except StopIteration:
dataiterator = iter(dataloader)
input_data = next(dataiterator)
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
net_outputs = maskRCNN(**input_data)
training_stats.UpdateIterStats(net_outputs, inner_iter)
loss = net_outputs['total_loss']
loss.backward()
optimizer.step()
training_stats.IterToc()
training_stats.LogIterStats(step, lr, backbone_lr)
if (step + 1) % EVAL_PERIOD == 0 or (step
== cfg.SOLVER.MAX_ITER - 1):
logger.info('Validating model')
eval_model = maskRCNN.module
eval_model = mynn.DataParallel(
eval_model,
cpu_keywords=['im_info', 'roidb'],
device_ids=[0],
minibatch=True)
eval_model.eval()
all_results = run_eval_inference(eval_model,
val_roidb,
args,
val_dataset,
val_dataset_name,
val_proposal_file,
ind_range=None,
multi_gpu_testing=False,
check_expected_results=True)
csv_path = os.path.join(output_dir, 'eval.csv')
all_results = all_results[0]
generate_csv_file_from_det_obj(all_results, csv_path,
obj_categories, prd_categories,
obj_freq_dict, prd_freq_dict)
overall_metrics, per_class_metrics = get_metrics_from_csv(
csv_path)
obj_acc = per_class_metrics[(csv_path, 'obj', 'top1')]
sbj_acc = per_class_metrics[(csv_path, 'sbj', 'top1')]
prd_acc = per_class_metrics[(csv_path, 'rel', 'top1')]
avg_obj_sbj = (obj_acc + sbj_acc) / 2.0
avg_acc = (prd_acc + avg_obj_sbj) / 2.0
best = json.load(open(os.path.join(ckpt_dir, 'best.json')))
if avg_acc > best['avg_per_class_acc']:
print('Found new best validation accuracy at {:2.2f}%'.
format(avg_acc))
print('Saving best model..')
best['avg_per_class_acc'] = avg_acc
best['iteration'] = step
best['per_class_metrics'] = {
'obj_top1':
per_class_metrics[(csv_path, 'obj', 'top1')],
'sbj_top1':
per_class_metrics[(csv_path, 'sbj', 'top1')],
'prd_top1':
per_class_metrics[(csv_path, 'rel', 'top1')]
}
best['overall_metrics'] = {
'obj_top1': overall_metrics[(csv_path, 'obj', 'top1')],
'sbj_top1': overall_metrics[(csv_path, 'sbj', 'top1')],
'prd_top1': overall_metrics[(csv_path, 'rel', 'top1')]
}
save_best_ckpt(output_dir, args, step, train_size,
maskRCNN, optimizer)
json.dump(best,
open(os.path.join(ckpt_dir, 'best.json'), 'w'))
if (step + 1) % CHECKPOINT_PERIOD == 0:
print('Saving Checkpoint..')
save_ckpt(output_dir, args, step, train_size, maskRCNN,
optimizer)
# ---- Training ends ----
# Save last checkpoint
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
except Exception as e:
del dataiterator
logger.info('Save ckpt on exception ...')
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
logger.info('Save ckpt done.')
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
if args.use_tfboard and not args.no_save:
tblogger.close()
cv2.imwrite(
osp.join(args.out_path, args.vid_name + '_vis',
'%08d.jpg' % (i + 1)), vis_im)
if __name__ == '__main__':
workspace.GlobalInit(['caffe2', '--caffe2_log_level=0'])
args = parse_args()
if args.out_path == None:
args.out_path = args.video_path
args.vid_name = args.video_path.split('/')[-1].split('.')[0]
utils.c2.import_custom_ops()
utils.c2.import_detectron_ops()
utils.c2.import_contrib_ops()
if args.cfg_file is not None:
cfg_from_file(args.cfg_file)
if args.opts is not None:
cfg_from_list(args.opts)
assert_and_infer_cfg()
if osp.exists(osp.join(args.out_path, args.vid_name + '_vis')):
shutil.rmtree(osp.join(args.out_path, args.vid_name + '_vis'))
os.makedirs(osp.join(args.out_path, args.vid_name + '_vis'))
num_images = _read_video(args)
gpu_dev = core.DeviceOption(caffe2_pb2.CUDA, cfg.ROOT_GPU_ID)
name_scope = 'gpu_{}'.format(cfg.ROOT_GPU_ID)
main(name_scope, gpu_dev, num_images, args)
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.input_video_path
assert args.output_path
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
s = VideoInputStream(args.input_video_path)
frame_id = 0
frame_detections = {}
for im in s:
print('frame_id', frame_id)
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, _ = im_detect_raw_masks(maskRCNN,
im,
timers=timers)
if cls_segms is not None:
print(im.shape, len(cls_boxes), len(cls_segms))
boxes, class_ids, scores, masks = \
process_detections(cls_boxes, cls_segms, s.height, s.width)
frame_detections[frame_id] = [boxes, class_ids, scores, masks]
frame_id = frame_id + 1
np.save(args.output_path, frame_detections)
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
# Pedestrian or Background
class Dataset:
def __init__(self):
self.classes = {0: 'background', 1: 'pedestrian'}
dataset = Dataset()
cfg.MODEL.NUM_CLASSES = 2
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
if not os.path.exists(osp.join(args.output_dir, args.model_name)):
os.makedirs(osp.join(args.output_dir, args.model_name))
for i in xrange(num_images):
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
osp.join(args.output_dir, args.model_name),
cls_boxes,
cls_segms,
cls_keyps,
dataset=dataset,
box_alpha=0.5,
show_class=True,
thresh=0.9,
kp_thresh=2,
ext='png')