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