if args.use_tb:
num_cpu = 1
tf_config = tf.ConfigProto(
inter_op_parallelism_threads=num_cpu,
intra_op_parallelism_threads=num_cpu,
gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.005))
sess = tf.Session(config=tf_config)
summary = Summary(sess, log_dir)
tb = logger.Logger(
log_dir, output_formats=[logger.TensorBoardOutputFormat(log_dir)])
# initilize the network here.
print(args.model)
if args.model == 'oicr':
OICR = vgg16_oicr(imdb.classes,
pretrained=True,
class_agnostic=args.class_agnostic,
summary=summary)
else:
raise Exception("Model does not exist")
OICR.create_architecture()
lr = cfg.TRAIN.LEARNING_RATE
lr = args.lr
#tr_momentum = cfg.TRAIN.MOMENTUM
#tr_momentum = args.momentum
weight_decay = cfg.TRAIN.WEIGHT_DECAY
param_groups = OICR.groups
if args.cuda:
imdb.evaluate_detections(None, output_dir_map, args.restore)
print('Evaluating CorLoc')
imdb.evaluate_discovery(None, output_dir_corloc, args.restore)
exit()
input_dir = args.load_dir + "/" + args.net + "/" + args.dataset
if not os.path.exists(input_dir):
raise Exception(
'There is no input directory for loading network from ' +
input_dir)
load_name = os.path.join(input_dir, '{:06d}.pth'.format(args.checkpoint))
# initilize the network here.
if args.model == 'oicr':
OICR = vgg16_oicr(imdb.classes,
pretrained=False,
class_agnostic=args.class_agnostic)
else:
raise Exception("Model does not exist")
OICR.create_architecture()
print("load checkpoint %s" % (load_name))
checkpoint = torch.load(load_name)
OICR.load_state_dict(checkpoint['model'])
if 'pooling_mode' in checkpoint.keys():
cfg.POOLING_MODE = checkpoint['pooling_mode']
print('load model successfully!')
if args.cuda: