layers.append(m_) ch.append(c2) return nn.Sequential(*layers), sorted(save), sorted(out_list) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') opt = parser.parse_args() opt.cfg = check_file(opt.cfg) # check file set_logging() device = select_device(opt.device) # Create model model = Model(opt.cfg).to(device) model.train() # Profile # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) # y = model(img, profile=True) # ONNX export # model.model[-1].export = True # torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11)
opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1 set_logging(opt.global_rank) #setup_seed(2020) # Resume if opt.resume: last = get_latest_run( ) if opt.resume == 'get_last' else opt.resume # resume from most recent run if last and not opt.weights: logger.info(f'Resuming training from {last}') opt.weights = last if opt.resume and not opt.weights else opt.weights if opt.global_rank in [-1, 0]: check_git_status() #opt.hyp = opt.hyp or ('data/hyp.finetune.yaml' if opt.weights else 'data/hyp.scratch.yaml') opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file( opt.cfg), check_file(opt.hyp) # check files assert len(opt.cfg) or len( opt.weights), 'either --cfg or --weights must be specified' opt.img_size.extend( [opt.img_size[-1]] * (2 - len(opt.img_size))) # extend to 2 sizes (train, test) device = select_device(opt.device, batch_size=opt.batch_size) # DDP mode if opt.local_rank != -1: assert torch.cuda.device_count() > opt.local_rank torch.cuda.set_device(opt.local_rank) device = torch.device('cuda', opt.local_rank) dist.init_process_group(backend='nccl',
parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--merge', action='store_true', help='use Merge NMS') parser.add_argument('--verbose', action='store_true', help='report mAP by class') parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') opt = parser.parse_args() opt.save_json |= opt.data.endswith('coco.yaml') opt.data = check_file(opt.data) # check file print(opt) if opt.task in ['val', 'test']: # run normally test(opt.data, opt.weights, opt.batch_size, opt.img_size, opt.conf_thres, opt.iou_thres, opt.save_json, opt.single_cls, opt.augment, opt.verbose) elif opt.task == 'study': # run over a range of settings and save/plot for weights in [ 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', 'yolov3-spp.pt' ]: f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem ) # filename to save to x = list(range(352, 832, 64)) # x axis