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)