def main():
args = parse_args()
logger = log.get_logger(args.log)
args.logger = logger
logger.info('*' * 80)
logger.info('the args are the below')
logger.info('*' * 80)
for x in args.__dict__:
logger.info(x + ',' + str(args.__dict__[x]))
logger.info(cfg.config[args.dataset])
logger.info('*' * 80)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if not os.path.exists(args.param_dir):
os.mkdir(args.param_dir)
torch.manual_seed(long(time.time()))
model = ablation.BDCN(pretrain=args.pretrain,
logger=logger,
ms=args.ms,
block=args.block,
bdcn=not args.no_bdcn,
direction=args.dir,
k=args.num_conv,
rate=args.rate)
if args.complete_pretrain:
model.load_state_dict(torch.load(args.complete_pretrain))
logger.info(model)
train(model, args)
def main():
args = parse_args()
logger = log.get_logger(args.log)
args.logger = logger
logger.info('*'*80)
logger.info('the args are the below')
logger.info('*'*80)
for x in args.__dict__:
logger.info(x+','+str(args.__dict__[x]))
logger.info(cfg.config[args.dataset])
logger.info('*'*80)
# os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
device = torch.device('cpu' if torch.cuda.device_count() == 0
else 'cuda')
if not os.path.exists(args.param_dir):
os.mkdir(args.param_dir)
torch.manual_seed(time.time())
model = ablation.BDCN(pretrain=args.pretrain, logger=logger,
ms=args.ms, block=args.block, bdcn=not args.no_bdcn, direction=args.dir,
k=args.num_conv, rate=args.rate).to(device)
res_model = count_parameters(model)
print('Number of parameters> ',res_model)
if args.complete_pretrain:
model.load_state_dict(torch.load(args.complete_pretrain))
logger.info(model)
train(model, args, devi= device)
def main():
import time
print time.localtime()
args = parse_args()
args.bdcn = not args.no_bdcn
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
model = ablation.BDCN(ms=args.ms, block=args.block, bdcn=not args.no_bdcn,
direction=args.dir, k=args.num_conv, rate=args.rate)
model.load_state_dict(torch.load('%s' % (args.model)))
test(model, args)
def main():
import time
print(time.localtime())
args = parse_args()
args.bdcn = not args.no_bdcn
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
device = torch.device('cpu' if torch.cuda.device_count() == 0 else 'cuda')
model = ablation.BDCN(ms=args.ms,
block=args.block,
bdcn=not args.no_bdcn,
direction=args.dir,
k=args.num_conv,
rate=args.rate).to(device)
chckpnt_dir = os.path.join('params',
args.train_data + '-B' + str(args.block),
args.model)
model.load_state_dict(torch.load(chckpnt_dir, map_location=device))
print('Successfuly checkpoint loaded ', chckpnt_dir)
test(model, args, device=device)