def main(): global args, logger, v_id args = parser.parse_args() cfg = load_config(args) init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') logger.info(args) # setup model if args.arch == 'Custom': from custom import Custom model = Custom(anchors=cfg['anchors']) else: parser.error('invalid architecture: {}'.format(args.arch)) if args.resume: assert isfile(args.resume), '{} is not a valid file'.format(args.resume) model = load_pretrain(model, args.resume) model.eval() device = torch.device('cuda' if (torch.cuda.is_available() and not args.cpu) else 'cpu') model = model.to(device) # setup dataset dataset = load_dataset(args.dataset) # VOS or VOT? if args.dataset in ['DAVIS','DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask: vos_enable = True # enable Mask output else: vos_enable = False total_lost = 0 # VOT iou_lists = [] # VOS speed_list = [] for v_id, video in enumerate(dataset.keys(), start=1): if args.video != '' and video != args.video: continue if vos_enable: iou_list, speed = track_vos(model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None, args.mask, args.refine, args.dataset in ['DAVIS2017', 'ytb_vos'], device=device) iou_lists.append(iou_list) else: lost, speed = track_vot(model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None, args.mask, args.refine, device=device) total_lost += lost speed_list.append(speed) # report final result if vos_enable: for thr, iou in zip(thrs, np.mean(np.concatenate(iou_lists), axis=0)): logger.info('Segmentation Threshold {:.2f} mIoU: {:.3f}'.format(thr, iou)) else: logger.info('Total Lost: {:d}'.format(total_lost)) logger.info('Mean Speed: {:.2f} FPS'.format(np.mean(speed_list)))
def __init__(self, config_path, model_path): args = TrackArgs() args.config = config_path args.resume = model_path cfg = load_config(args) if args.arch == 'Custom': from custom import Custom self.model = Custom(anchors=cfg['anchors']) else: parser.error('invalid architecture: {}'.format(args.arch)) if args.resume: assert isfile(args.resume), '{} is not a valid file'.format(args.resume) self.model = load_pretrain(self.model, args.resume) self.model.eval() self.device = torch.device('cuda' if (torch.cuda.is_available() and not args.cpu) else 'cpu') self.model = self.model.to(self.device) ################# Dangerous self.p = TrackerConfig() self.p.update(cfg['hp'] if 'hp' in cfg.keys() else None, self.model.anchors) self.p.renew() self.p.scales = self.model.anchors['scales'] self.p.ratios = self.model.anchors['ratios'] self.p.anchor_num = self.model.anchor_num self.p.anchor = generate_anchor(self.model.anchors, self.p.score_size) if self.p.windowing == 'cosine': self.window = np.outer(np.hanning(self.p.score_size), np.hanning(self.p.score_size)) elif self.p.windowing == 'uniform': self.window = np.ones((self.p.score_size, self.p.score_size)) self.window = np.tile(self.window.flatten(), self.p.anchor_num)
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() # args通过解析获得的 init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') # 实例化一个记录器 logger.info("\n" + collect_env_info()) logger.info(args) cfg = load_config(args) logger.info("config \n{}".format(json.dumps( cfg, indent=4))) # 转变成json格式的文件,缩进4格 if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # build dataset train_loader, val_loader = build_data_loader(cfg) if args.arch == 'Custom': from custom import Custom model = Custom(pretrain=True, anchors=cfg['anchors']) else: exit() logger.info(model) if args.pretrained: model = load_pretrain(model, args.pretrained) model = model.cuda() # 模型转移到GPU上 dist_model = torch.nn.DataParallel( model, list(range(torch.cuda.device_count()))).cuda() # 多GPU训练 if args.resume and args.start_epoch != 0: # 这是在干啥?蒙蔽了!!!!! model.features.unfix((args.start_epoch - 1) / args.epochs) optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch) # 如何构建优化器和学习策略??? # optionally resume from a checkpoint if args.resume: assert os.path.isfile(args.resume), '{} is not a valid file'.format( args.resume) model, optimizer, args.start_epoch, best_acc, arch = restore_from( model, optimizer, args.resume) dist_model = torch.nn.DataParallel( model, list(range(torch.cuda.device_count()))).cuda() logger.info(lr_scheduler) logger.info('model prepare done') train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg)
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() init_log('global', logging.INFO) # 返回一个logger对象,logging_INFO是日志的等级 if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') # 获取上面初始化的logger对象 logger.info("\n" + collect_env_info()) logger.info(args) cfg = load_config(args) # 返回修改后的配置文件对象 logger.info("config \n{}".format(json.dumps(cfg, indent=4))) #json.loads()是将str转化成dict格式,json.dumps()是将dict转化成str格式。 if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # build dataset train_loader, val_loader = build_data_loader(cfg) if args.arch == 'Custom': from custom import Custom model = Custom(pretrain=True, anchors=cfg['anchors']) else: exit() logger.info(model) if args.pretrained: model = load_pretrain(model, args.pretrained) model = model.cuda() dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda() if args.resume and args.start_epoch != 0: model.features.unfix((args.start_epoch - 1) / args.epochs) optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch) # optionally resume from a checkpoint if args.resume: assert os.path.isfile(args.resume), '{} is not a valid file'.format(args.resume) model, optimizer, args.start_epoch, best_acc, arch = restore_from(model, optimizer, args.resume) dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda() logger.info(lr_scheduler) logger.info('model prepare done') train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg)
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() args = args_process(args) init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') logger.info("\n" + collect_env_info()) logger.info(args) cfg = load_config(args) logger.info("config \n{}".format(json.dumps(cfg, indent=4))) # build dataset train_loader, val_loader = build_data_loader(cfg) args.img_size = int(cfg['train_datasets']['search_size']) args.nms_threshold = float(cfg['train_datasets']['RPN_NMS']) if args.arch == 'Custom': from custom import Custom model = Custom(pretrain=True, opts=args, anchors=train_loader.dataset.anchors) else: exit() logger.info(model) if args.pretrained: model = load_pretrain(model, args.pretrained) else: raise Exception("Pretrained weights must be loaded!") model = model.cuda() dist_model = torch.nn.DataParallel(model, list(range( torch.cuda.device_count()))).cuda() logger.info('model prepare done') logger = logging.getLogger('global') val_avg = AverageMeter() validation(val_loader, dist_model, cfg, val_avg)
def __init__(self, args): super(PatchTrainer, self).__init__() # Setup device self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') torch.backends.cudnn.benchmark = True # Setup tracker cfg cfg = load_config(args) p = TrackerConfig() p.renew() self.p = p # Setup tracker siammask = Tracker(p=p, anchors=cfg['anchors']) if args.resume: assert isfile(args.resume), 'Please download {} first.'.format(args.resume) siammask = load_pretrain(siammask, args.resume) siammask.eval().to(self.device) self.model = siammask
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') logger.info("\n" + collect_env_info()) logger.info(args) cfg = load_config(args) logger.info("config \n{}".format(json.dumps(cfg, indent=4))) if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # build dataset train_loader, val_loader = build_data_loader(cfg)
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() args = args_process(args) init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') logger.info("\n" + collect_env_info()) logger.info(args) cfg = load_config(args) logger.info("config \n{}".format(json.dumps(cfg, indent=4))) if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # build dataset train_loader, val_loader = build_data_loader(cfg) args.img_size = int(cfg['train_datasets']['search_size']) args.nms_threshold = float(cfg['train_datasets']['RPN_NMS']) if args.arch == 'Custom': from custom import Custom model = Custom(pretrain=True, opts=args, anchors=train_loader.dataset.anchors) else: exit() logger.info(model) if args.pretrained: model = load_pretrain(model, args.pretrained) model = model.cuda() dist_model = torch.nn.DataParallel(model, list(range( torch.cuda.device_count()))).cuda() if args.resume and args.start_epoch != 0: model.features.unfix((args.start_epoch - 1) / args.epochs) optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch) # optionally resume from a checkpoint if args.resume: assert os.path.isfile(args.resume), '{} is not a valid file'.format( args.resume) model, optimizer, args.start_epoch, best_acc, arch = restore_from( model, optimizer, args.resume) dist_model = torch.nn.DataParallel( model, list(range(torch.cuda.device_count()))).cuda() logger.info(lr_scheduler) logger.info('model prepare done') global cur_lr if not os.path.exists(args.save_dir): # makedir/save model os.makedirs(args.save_dir) num_per_epoch = len(train_loader.dataset) // args.batch num_per_epoch_val = len(val_loader.dataset) // args.batch for epoch in range(args.start_epoch, args.epochs): lr_scheduler.step(epoch) cur_lr = lr_scheduler.get_cur_lr() logger = logging.getLogger('global') train_avg = AverageMeter() val_avg = AverageMeter() if dist_model.module.features.unfix(epoch / args.epochs): logger.info('unfix part model.') optimizer, lr_scheduler = build_opt_lr(dist_model.module, cfg, args, epoch) train(train_loader, dist_model, optimizer, lr_scheduler, epoch, cfg, train_avg, num_per_epoch) if dist_model.module.features.unfix(epoch / args.epochs): logger.info('unfix part model.') optimizer, lr_scheduler = build_opt_lr(dist_model.module, cfg, args, epoch) if (epoch + 1) % args.save_freq == 0: save_checkpoint( { 'epoch': epoch, 'arch': args.arch, 'state_dict': dist_model.module.state_dict(), 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), 'anchor_cfg': cfg['anchors'] }, False, os.path.join(args.save_dir, 'checkpoint_e%d.pth' % (epoch)), os.path.join(args.save_dir, 'best.pth')) validation(val_loader, dist_model, epoch, cfg, val_avg, num_per_epoch_val)
default='../../data/tennis', help='datasets') parser.add_argument('--gt_file', default=None, type=str, help='ground truth txt file') parser.add_argument('--cpu', action='store_true', help='cpu mode') args = parser.parse_args() if __name__ == '__main__': # Setup device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') torch.backends.cudnn.benchmark = True # Setup Model cfg = load_config(args) siammask = Custom_(anchors=cfg['anchors']) if args.resume: assert isfile(args.resume), 'Please download {} first.'.format( args.resume) siammask = load_pretrain(siammask, args.resume) siammask.eval().to(device) # Parse Image file img_files = sorted(glob.glob(join(args.base_path, '*.jp*'))) ims = [cv2.imread(imf) for imf in img_files] # Select ROI cv2.namedWindow("SiamMask", cv2.WND_PROP_FULLSCREEN) # cv2.setWindowProperty("SiamMask", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
def main(): global args, logger, v_id args = parser.parse_args() cfg = load_config(args) init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') logger.info(args) # setup model if args.arch == 'Custom': from custom import Custom model = Custom(anchors=cfg['anchors']) else: parser.error('invalid architecture: {}'.format(args.arch)) if args.resume: assert isfile(args.resume), '{} is not a valid file'.format( args.resume) model = load_pretrain(model, args.resume) model.eval() device = torch.device('cuda' if ( torch.cuda.is_available() and not args.cpu) else 'cpu') model = model.to(device) # setup dataset dataset = load_dataset(args.dataset, args.dir_type) # VOS or VOT? if args.dataset in ['DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask: vos_enable = True # enable Mask output else: vos_enable = False # total_lost = 0 # VOT # iou_lists = [] # VOS # speed_list = [] for v_id, video in enumerate(dataset.keys(), start=1): if args.video != '' and video != args.video: continue if vos_enable: iou_list, speed = track_vos( model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None, args.mask, args.refine, args.dataset in ['DAVIS2017', 'ytb_vos'], device=device) # iou_lists.append(iou_list) else: lost, speed = track_vot(model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None, args.mask, args.refine, device=device) total_lost += lost
def process_vedio(vedio_path, initRect): """ 视频处理 :param vedio_path:视频路径 :param initRect: 跟踪目标的初始位置 :return: """ # 1. 设置设备信息 Setup device # 有GPU时选择GPU,否则使用CPU device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 默认优化运行效率 torch.backends.cudnn.benchmark = True # 2. 模型设置 Setup Model # 2.1 将命令行参数解析出来 cfg = load_config(args) # 2.2 custom是构建的网络,否则引用model中的网络结构 from custom import Custom siammask = Custom(anchors=cfg['anchors']) # 2.3 判断是否存在模型的权重文件 if args.resume: assert isfile(args.resume), 'Please download {} first.'.format( args.resume) siammask = load_pretrain(siammask, args.resume) # 在运行推断前,需要调用 model.eval() 函数,以将 dropout 层 和 batch normalization 层设置为评估模式(非训练模式). # to(device)将张量复制到GPU上,之后的计算将在GPU上运行 siammask.eval().to(device) # 首帧跟踪目标的位置 x, y, w, h = initRect print(x) VeryBig = 999999999 # 用于将视频框调整到最大 Cap = cv2.VideoCapture(vedio_path) # 设置读取摄像头 ret, frame = Cap.read() # 读取帧 ims = [frame] # 把frame放入列表格式的frame, 因为原文是将每帧图片放入列表 im = frame f = 0 target_pos = np.array([x + w / 2, y + h / 2]) target_sz = np.array([w, h]) state = siamese_init(im, target_pos, target_sz, siammask, cfg['hp']) # init tracker" middlepath = "../data/middle.mp4" outpath = "../data/output.mp4" vediowriter = cv2.VideoWriter(middlepath, cv2.VideoWriter_fourcc('M', 'P', '4', 'V'), 10, (320, 240)) while (True): tic = cv2.getTickCount() ret, im = Cap.read() # 逐个提取frame if (ret == False): break state = siamese_track(state, im, mask_enable=True, refine_enable=True) # track location = state['ploygon'].flatten() mask = state['mask'] > state['p'].seg_thr im[:, :, 2] = (mask > 0) * 255 + (mask == 0) * im[:, :, 2] cv2.polylines(im, [np.int0(location).reshape((-1, 1, 2))], True, (0, 255, 0), 3) vediowriter.write(im) cv2.imshow('SiamMask', im) key = cv2.waitKey(1) if key > 0: break f = f + 1 vediowriter.release() return
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) print("Init logger") logger = logging.getLogger('global') print(44) #logger.info("\n" + collect_env_info()) print(99) logger.info(args) cfg = load_config(args) logger.info("config \n{}".format(json.dumps(cfg, indent=4))) print(2) if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # build dataset train_loader, val_loader = build_data_loader(cfg) print(3) path = "/usr4/alg504/cliao25/siammask/experiments/siammask_base/snapshot/checkpoint_e{}.pth" for epoch in range(1,21): if args.arch == 'Custom': from custom import Custom model = Custom(pretrain=True, anchors=cfg['anchors']) else: exit() print(4) if args.pretrained: model = load_pretrain(model, args.pretrained) model = model.cuda() #model.features.unfix((epoch - 1) / 20) optimizer, lr_scheduler = build_opt_lr(model, cfg, args, epoch) filepath = path.format(epoch) assert os.path.isfile(filepath) model, _, _, _, _ = restore_from(model, optimizer, filepath) #model = load_pretrain(model, filepath) model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda() model.train() device = torch.device('cuda') model = model.to(device) valid(val_loader, model, cfg) print("Done")
def main(): # 获取命令行参数信息 global args, logger, v_id args = parser.parse_args() # 获取配置文件中配置信息:主要包括网络结构,超参数等 cfg = load_config(args) # 初始化logxi信息,并将日志信息输入到磁盘文件中 init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) # 将相关的配置信息输入到日志文件中 logger = logging.getLogger('global') logger.info(args) # setup model # 加载网络模型架构 if args.arch == 'Custom': from custom import Custom model = Custom(anchors=cfg['anchors']) else: parser.error('invalid architecture: {}'.format(args.arch)) # 加载网络模型参数 if args.resume: assert isfile(args.resume), '{} is not a valid file'.format( args.resume) model = load_pretrain(model, args.resume) # 使用评估模式,将drop等激活 model.eval() # 硬件信息 device = torch.device('cuda' if ( torch.cuda.is_available() and not args.cpu) else 'cpu') model = model.to(device) # 加载数据集 setup dataset dataset = load_dataset(args.dataset) # 这三种数据支持掩膜 VOS or VOT? if args.dataset in ['DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask: vos_enable = True # enable Mask output else: vos_enable = False total_lost = 0 # VOT iou_lists = [] # VOS speed_list = [] # 对数据进行处理 for v_id, video in enumerate(dataset.keys(), start=1): if args.video != '' and video != args.video: continue # true 调用track_vos if vos_enable: # 如测试数据是['DAVIS2017', 'ytb_vos']时,会开启多目标跟踪 iou_list, speed = track_vos( model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None, args.mask, args.refine, args.dataset in ['DAVIS2017', 'ytb_vos'], device=device) iou_lists.append(iou_list) # False 调用track_vot else: lost, speed = track_vot(model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None, args.mask, args.refine, device=device) total_lost += lost speed_list.append(speed) # report final result if vos_enable: for thr, iou in zip(thrs, np.mean(np.concatenate(iou_lists), axis=0)): logger.info('Segmentation Threshold {:.2f} mIoU: {:.3f}'.format( thr, iou)) else: logger.info('Total Lost: {:d}'.format(total_lost)) logger.info('Mean Speed: {:.2f} FPS'.format(np.mean(speed_list)))
def main(): # args.base_path = base_path args.resume = "../SiamMask/experiments/siammask_sharp/SiamMask_DAVIS.pth" args.config = "../SiamMask/experiments/siammask_sharp/config_davis.json" print(join(args.base_path, 'groundtruth_rect.txt')) # Setup device device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') torch.backends.cudnn.benchmark = True # Setup Model cfg = load_config(args) p = TrackerConfig() p.renew() siammask = Tracker(p=p, anchors=cfg['anchors']) if args.resume: assert isfile(args.resume), 'Please download {} first.'.format(args.resume) siammask = load_pretrain(siammask, args.resume) siammask.eval().to(device) # Parse Image file img_files = sorted(glob.glob(join(join(args.base_path, 'imgs'), '*.jp*'))) ims = [cv2.imread(imf) for imf in img_files] # Select ROI cv2.namedWindow("SiamMask", cv2.WND_PROP_FULLSCREEN) try: init_rect = cv2.selectROI('SiamMask', ims[0], False, False) gts = None x, y, w, h = init_rect except: exit() file1 = open(join(args.base_path, 'groundtruth_rect.txt'), 'w') file1.write('{0:d},{1:d},{2:d},{3:d}\n'.format(x, y, w, h)) toc = 0 for f, im in enumerate(ims): tic = cv2.getTickCount() if f == 0: # init target_pos = np.array([x + w / 2, y + h / 2]) target_sz = np.array([w, h]) state = tracker_init(im, target_pos, target_sz, siammask, device=device) # init tracker state['gts'] = gts state['device'] = device elif f > 0: # tracking state = tracker_track(state, im, siammask, device=device) # track target_pos, target_sz =state['target_pos'], state['target_sz'] x, y = (target_pos - target_sz/2).astype(int) x2, y2 = (target_pos + target_sz/2).astype(int) cv2.rectangle(im, (x, y), (x2, y2), (0, 255, 0), 4) cv2.imshow('SiamMask', im) key = cv2.waitKey(1) if key == ord('q'): break file1.write('{0:d},{1:d},{2:d},{3:d}\n'.format(x, y, x2-x, y2-y)) toc += cv2.getTickCount() - tic file1.close() toc /= cv2.getTickFrequency() fps = f / toc print('SiamMask Time: {:02.1f}s Speed: {:3.1f}fps (with visulization!)'.format(toc, fps))
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() args = args_process(args) init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') logger.info("\n" + collect_env_info()) logger.info(args) cfg = load_config(args) logger.info("config \n{}".format(json.dumps(cfg, indent=4))) if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # build dataset train_loader, val_loader = build_data_loader(cfg) args.img_size = int(cfg['train_datasets']['search_size']) args.nms_threshold = float(cfg['train_datasets']['RPN_NMS']) if args.arch == 'Custom': from custom import Custom model = Custom(pretrain=True, opts=args, anchors=train_loader.dataset.anchors) else: exit() logger.info(model) if args.pretrained: model = load_pretrain(model, args.pretrained) model = model.cuda() dist_model = torch.nn.DataParallel(model, list(range( torch.cuda.device_count()))).cuda() if args.resume and args.start_epoch != 0: model.features.unfix((args.start_epoch - 1) / args.epochs) optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch) # optionally resume from a checkpoint if args.resume: assert os.path.isfile(args.resume), '{} is not a valid file'.format( args.resume) model, optimizer, args.start_epoch, best_acc, arch = restore_from( model, optimizer, args.resume) dist_model = torch.nn.DataParallel( model, list(range(torch.cuda.device_count()))).cuda() logger.info(lr_scheduler) logger.info('model prepare done') train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg)
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') logger.info(args) cfg = load_config(args) logger.info("config \n{}".format(json.dumps(cfg, indent=4))) logger.info("\n" + collect_env_info()) if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # build dataset train_loader, val_loader = build_data_loader(cfg) if args.arch == 'Custom': from custom import Custom model = Custom(pretrain=True, anchors=cfg['anchors']) else: model = models.__dict__[args.arch](anchors=cfg['anchors']) logger.info(model) if args.pretrained: model = load_pretrain(model, args.pretrained) model = model.cuda() dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda() if args.resume and args.start_epoch != 0: model.features.unfix((args.start_epoch - 1) / args.epochs) optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch) logger.info(lr_scheduler) # optionally resume from a checkpoint if args.resume: assert os.path.isfile(args.resume), '{} is not a valid file'.format(args.resume) model, optimizer, args.start_epoch, best_acc, arch = restore_from(model, optimizer, args.resume) dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda() epoch = args.start_epoch if dist_model.module.features.unfix(epoch/args.epochs): logger.info('unfix part model.') optimizer, lr_scheduler = build_opt_lr(dist_model.module, cfg, args, epoch) lr_scheduler.step(epoch) cur_lr = lr_scheduler.get_cur_lr() logger.info('epoch:{} resume lr {}'.format(epoch, cur_lr)) logger.info('model prepare done') train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg)
def main(): init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) params = {'penalty_k': args.penalty_k, 'window_influence': args.window_influence, 'lr': args.lr, 'instance_size': args.search_region} num_search = len(params['penalty_k']) * len(params['window_influence']) * \ len(params['lr']) * len(params['instance_size']) print(params) print(num_search) cfg = load_config(args) if args.arch == 'Custom': from custom import Custom model = Custom(anchors=cfg['anchors']) else: model = models.__dict__[args.arch](anchors=cfg['anchors']) if args.resume: assert isfile(args.resume), '{} is not a valid file'.format(args.resume) model = load_pretrain(model, args.resume) model.eval() model = model.to(device) default_hp = cfg.get('hp', {}) p = dict() p['network'] = model p['network_name'] = args.arch+'_'+args.resume.split('/')[-1].split('.')[0] p['dataset'] = args.dataset global ims, gt, image_files dataset_info = load_dataset(args.dataset) videos = list(dataset_info.keys()) np.random.shuffle(videos) for video in videos: print(video) if isfile('finish.flag'): return p['video'] = video ims = None image_files = dataset_info[video]['image_files'] gt = dataset_info[video]['gt'] np.random.shuffle(params['penalty_k']) np.random.shuffle(params['window_influence']) np.random.shuffle(params['lr']) for penalty_k in params['penalty_k']: for window_influence in params['window_influence']: for lr in params['lr']: for instance_size in params['instance_size']: p['hp'] = default_hp.copy() p['hp'].update({'penalty_k':penalty_k, 'window_influence':window_influence, 'lr':lr, 'instance_size': instance_size, }) tune(p)
def main(): """ 基础网络的训练 :return: """ global args, best_acc, tb_writer, logger args = parser.parse_args() # 初始化日志信息 init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) # 获取log信息 logger = logging.getLogger('global') logger.info("\n" + collect_env_info()) logger.info(args) # 获取配置信息 cfg = load_config(args) logger.info("config \n{}".format(json.dumps(cfg, indent=4))) if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # 构建数据集 train_loader, val_loader = build_data_loader(cfg) # 加载训练网络 if args.arch == 'Custom': from custom import Custom model = Custom(pretrain=True, anchors=cfg['anchors']) else: exit() logger.info(model) # 加载预训练网络 if args.pretrained: model = load_pretrain(model, args.pretrained) # GPU版本 # model = model.cuda() # dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda() # 网络模型 dist_model = torch.nn.DataParallel(model) # 模型参数的更新比例 if args.resume and args.start_epoch != 0: model.features.unfix((args.start_epoch - 1) / args.epochs) # 获取优化器和学习率的更新策略 optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch) # optionally resume from a checkpoint 加载模型 if args.resume: assert os.path.isfile(args.resume), '{} is not a valid file'.format( args.resume) model, optimizer, args.start_epoch, best_acc, arch = restore_from( model, optimizer, args.resume) # GPU # dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda() dist_model = torch.nn.DataParallel(model) logger.info(lr_scheduler) logger.info('model prepare done') # 模型训练 train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg)
def main(): global args, best_acc, tb_writer, logger args = parser.parse_args() init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) logger = logging.getLogger('global') logger.info("\n" + collect_env_info()) logger.info(args) cfg = load_config(args) logger.info("config \n{}".format(json.dumps(cfg, indent=4))) if args.log_dir: tb_writer = SummaryWriter(args.log_dir) else: tb_writer = Dummy() # build dataset train_loader, val_loader = build_data_loader(cfg) if args.arch == 'Custom': model = Custom(anchors=cfg['anchors']) elif args.arch == 'Custom_Sky': model = Custom_Sky(anchors=cfg['anchors']) else: exit() logger.info(model) if args.pretrained: model = load_pretrain(model, args.pretrained) # print(summary(model=model, input_size=(3, 511, 511), batch_size=1)) model = model.cuda() dist_model = torch.nn.DataParallel(model, list(range( torch.cuda.device_count()))).cuda() if args.resume and args.start_epoch != 0: model.features.unfix((args.start_epoch - 1) / args.epochs) optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch) # optionally resume from a checkpoint if args.resume: print(args.resume) assert os.path.isfile(args.resume), '{} is not a valid file'.format( args.resume) model, optimizer, args.start_epoch, best_acc, arch = restore_from( model, optimizer, args.resume) dist_model = torch.nn.DataParallel( model, list(range(torch.cuda.device_count()))).cuda() logger.info(lr_scheduler) logger.info('model prepare done') train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg)
def main(): global args, logger, v_id #全局变量 args = parser.parse_args() #args是test.py文件运行时,接受的参数 cfg = load_config(args) #加载 JSON 配置文件并设置args.arch的值。 print(cfg) init_log('global', logging.INFO) if args.log != "": add_file_handler('global', args.log, logging.INFO) #add_file_handler 创建一个记录器并绑定文件句柄。 logger = logging.getLogger('global') logger.info(args) # setup model Custom 为论文实现的网络。如果不是“Custom”,加载 models 下指定的结构。 if args.arch == 'Custom': #args.arch参数,预训练模型的结构,命令行不给的话,默认为' ', from custom import Custom model = Custom(anchors=cfg['anchors'] ) #cfg是从config_vot.json的到的数据,所以跟踪时用的model.anchors字典中的数据 else: parser.error('invalid architecture: {}'.format(args.arch)) if args.resume: #给了args.resume,如果args.resume不是文件,报错, assert isfile(args.resume), '{} is not a valid file'.format( args.resume) model = load_pretrain( model, args.resume) #args.resume是文件load_pretrain ,能够处理网络之间的不一致 model.eval() device = torch.device('cuda' if ( torch.cuda.is_available() and not args.cpu) else 'cpu') model = model.to(device) # setup dataset,字典 dataset = load_dataset( args.dataset) #load_dataset 能够加载 VOT、DAVIS、ytb_vos 三种数据集。 #仅以上三种数据源支持掩膜输出。 # VOS or VOT? if args.dataset in ['DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask: vos_enable = True # enable Mask output ,使用掩膜输出 else: vos_enable = False total_lost = 0 # VOT 跟踪任务有损失函数 iou_lists = [] # VOS 分割任务 speed_list = [] #v_id视频索引从1起,video是视频名字 for v_id, video in enumerate(dataset.keys(), start=1): if v_id == 2: exit() if args.video != '' and video != args.video: #不成立,args.video默认是' ' continue if vos_enable: #分割任务,,,,分割任务和跟踪任务只能选一个 iou_list, speed = track_vos( model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None, args.mask, args.refine, args.dataset in ['DAVIS2017', 'ytb_vos'], device=device) iou_lists.append(iou_list) #iou_list是什么类型的数据??? else: #跟踪任务 lost, speed = track_vot(model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None, args.mask, args.refine, device=device) total_lost += lost speed_list.append(speed) # report final result记录最终结果 if vos_enable: #如果进行的是分割任务 for thr, iou in zip(thrs, np.mean(np.concatenate(iou_lists), axis=0)): logger.info('Segmentation Threshold {:.2f} mIoU: {:.3f}'.format( thr, iou)) else: logger.info('Total Lost: {:d}'.format(total_lost)) logger.info('Mean Speed: {:.2f} FPS'.format(np.mean(speed_list)))