def __init__(self, config): self.config = config self.exp_name = self.config.get('exp_name', None) if self.exp_name is None: self.exp_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") self.log_dir = osp.join(self.config['exp_dir'], self.exp_name, 'logs') self.pth_dir = osp.join(self.config['exp_dir'], self.exp_name, 'checkpoints') os.makedirs(self.log_dir, exist_ok=True) os.makedirs(self.pth_dir, exist_ok=True) self.writer = SummaryWriter(log_dir=self.log_dir) self.model = self._init_net() self.optimizer = self._init_optimizer() self.criterion = nn.CrossEntropyLoss().to(self.config['device']) self.train_loader, self.val_loader = self._init_dataloaders() pretrained_path = self.config.get('model_path', False) if pretrained_path: self.training_epoch, self.total_iter = load_checkpoint( pretrained_path, self.model, optimizer=self.optimizer) else: self.training_epoch = 0 self.total_iter = 0 self.epochs = self.config.get('epochs', int(1e5))
def __init__(self, config): self.config = config self.exp_name = self.config.get('exp_name', None) if self.exp_name is None: self.exp_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") self.res_dir = osp.join(self.config['exp_dir'], self.exp_name, 'results') os.makedirs(self.res_dir, exist_ok=True) self.model = self._init_net() self.inference_loader = self._init_dataloader() pretrained_path = self.config.get('model_path', False) if pretrained_path: load_checkpoint(pretrained_path, self.model) else: raise Exception( "model_path doesnt't exist in config. Please specify checkpoint path" )
def __init__(self, config): self.config = config self.exp_name = self.config.get("exp_name", None) if self.exp_name is None: self.exp_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") self.res_dir = osp.join(self.config["exp_dir"], self.exp_name, "results") os.makedirs(self.res_dir, exist_ok=True) self.model = self._init_net() self.pipeline_loader = self._init_dataloader() self.mapper = json.load(open(config["mapping_json"])) self.mapper = {j: i for i, j in self.mapper.items()} pretrained_path = self.config.get("model_path", False) if pretrained_path: load_checkpoint(pretrained_path, self.model) else: raise Exception( "model_path doesnt't exist in config. Please specify checkpoint path", )
def test_model(model, args): assert args.batch_size == 1 # get data list for tracking tracking_list_seq = [] tracking_list = [] batch_time = tu.AverageMeter() data_time = tu.AverageMeter() # resume from a checkpoint nu.load_checkpoint(model, args.resume, is_test=True) cudnn.benchmark = True dataset = Dataset(args.json_path, 'test', args.data_split, args.set == 'kitti', args.percent, args.is_tracking, args.is_normalizing, args.n_box_limit) print("Number of image to test: {}".format(dataset.__len__())) # Data loading code test_loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, drop_last=True) model.eval() end = time.time() for i, (image, box_info) in enumerate(iter(test_loader)): # measure data loading time data_time.update(time.time() - end) end = time.time() with torch.no_grad(): box_output, targets = model(image, box_info, args.device, 'test') batch_time.update(time.time() - end) rois_gt, \ dim_gt_, \ alpha_gt_, \ dep_gt_, \ cen_gt_, \ loc_gt_, \ ignore_, \ tid_gt = targets cam_calib = box_info['cam_calib'].cpu().data.numpy().reshape(3, 4) cam_rot = box_info['cam_rot'].cpu().data.numpy() cam_loc = box_info['cam_loc'].cpu().data.numpy() box_gt = rois_gt.cpu().data.numpy() box_pd = box_output['rois'].cpu().data.numpy() dim_gt = dim_gt_.cpu().data.numpy() dim_pd = box_output['dim'].cpu().data.numpy() alpha_gt = alpha_gt_.cpu().data.numpy() alpha_pd = nu.get_alpha(box_output['rot'].cpu().data.numpy()) depth_gt = dep_gt_.cpu().data.numpy() depth_pd = box_output['dep'].cpu().data.numpy() center_gt = cen_gt_.cpu().data.numpy() center_pd = box_output['cen'].cpu().data.numpy() loc_gt = loc_gt_.cpu().data.numpy() loc_pd = box_output['loc'].cpu().data.numpy() feature = F.normalize( F.avg_pool2d(box_output['feat'], (7, 7)).view(-1, 128)) # feature = box_output['feat'] feature_np = feature.cpu().data.numpy() tracking_list.append({ 'im_path': box_info['im_path'], 'endvid': box_info['endvid'].cpu().data.numpy(), 'rois_pd': box_pd, 'rois_gt': box_gt, 'feature': feature_np, 'dim_pd': dim_pd, 'alpha_pd': alpha_pd, 'depth_pd': depth_pd, 'center_pd': center_pd, 'loc_pd': loc_pd, 'dim_gt': dim_gt, 'alpha_gt': alpha_gt, 'depth_gt': depth_gt, 'center_gt': center_gt, 'loc_gt': loc_gt, 'cam_calib': cam_calib, 'cam_rot': cam_rot, 'cam_loc': cam_loc, 'ignore': ignore_.cpu().data.numpy(), 'tid_gt': tid_gt.cpu().data.numpy(), }) if box_info['endvid'].cpu().data.numpy().any() \ or i == len(test_loader): tracking_list_seq.append(tracking_list) tracking_list = [] if i % 100 == 0 and i != 0: print(i) end = time.time() if args.track_name is None: trk_name = os.path.join( cfg.OUTPUT_PATH, '{}_{}_{}_bdd_roipool_output.pkl'.format( args.session, str(args.start_epoch).zfill(3), args.set)) else: trk_name = os.path.join(cfg.OUTPUT_PATH, args.track_name) with open(trk_name, 'wb') as f: print("Saving {} with total {} sequences...".format( trk_name, len(tracking_list_seq))) pickle.dump(tracking_list_seq, f)
def run_training(model, args): if not os.path.isdir(cfg.CHECKPOINT_PATH): os.mkdir(cfg.CHECKPOINT_PATH) cudnn.benchmark = True if args.has_val: phases = ['train', 'val'] else: phases = ['train'] # Data loading code train_loader = { phase: DataLoader(Dataset(args.json_path, phase, phase, args.set == 'kitti', args.percent, args.is_tracking, args.is_normalizing, args.n_box_limit), batch_size=args.batch_size, shuffle=(phase == 'train'), num_workers=args.workers, pin_memory=True, drop_last=True) for phase in phases } # Optimizer if args.adaptBN: model_param = list() for m in model.modules(): if isinstance(m, nn.BatchNorm2d): model_param.append({'params': list(m.parameters())}) lr = 0.0 else: # model_param = model.parameters() model_param = filter(lambda p: p.requires_grad, model.parameters()) lr = args.lr if args.optim == 'sgd': optimizer = torch.optim.SGD(model_param, lr, momentum=args.momentum, nesterov=True, weight_decay=args.weight_decay) elif args.optim == 'adam': optimizer = torch.optim.Adam(model_param, lr, weight_decay=args.weight_decay, amsgrad=True) # optionally resume from a checkpoint if args.resume: nu.load_checkpoint(model, args.resume, optimizer=optimizer) # switch to train mode if args.adaptBN: model.eval() for m in model.modules(): if isinstance(m, nn.BatchNorm2d): m.train() for p in m.parameters(): p.requires_grad = False if args.use_tfboard: from tensorboardX import SummaryWriter logger = SummaryWriter("logs") else: logger = None for epoch in range(args.start_epoch, args.epochs): # Resume with normal lr adjust # not to over suppress the lr due to large epoch # if args.optim == 'sgd': nu.adjust_learning_rate(args, optimizer, epoch) # train for one epoch for phase in phases: if phase == 'train': model.train() # nu.freeze_model(model.module.base) # nu.freeze_model(model.module.rot) # nu.freeze_model(model.module.dep) # nu.freeze_model(model.module.dim) train_model(args, train_loader[phase], model, optimizer, epoch, phase, logger) acc = 0.0 else: model.eval() acc = val_model(args, train_loader[phase], model, epoch, phase, logger) # Save checkpoint nu.save_checkpoint( { 'epoch': epoch + 1, 'arch': args.arch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'best_score': acc, 'phase': phase, 'save_path': cfg.CHECKPOINT_PATH, }, args.set == 'kitti', args.session, args.check_freq) if args.use_tfboard: logger.close() print("Training finished!!")