def toonnx(args): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if not args.full_model: model = darknet.DarkNet().to(device) model_state_dict = model.state_dict() state_dict = torch.load(args.pytorch_model, map_location=device) state_dict = { k: v for k, v in state_dict.items() if k in model_state_dict } state_dict = collections.OrderedDict(state_dict) model_state_dict.update(state_dict) model.load_state_dict(model_state_dict) else: print('Warning: this function has not been tested yet!') model = torch.load(args.pytorch_model) dummy_input = torch.rand(1, 3, args.insize[0], args.insize[1], device=device) onnx.export(model, dummy_input, args.onnx_model, verbose=True, input_names=['data']) session = ort.InferenceSession(args.onnx_model) outputs = session.run(None, {'data': dummy_input.cpu().numpy()}) for i, output in enumerate(outputs): print('branch {} output size is {}'.format(i, output.shape))
def main(args): net = darknet.DarkNet(np.random.randint(0, 100, (12, 2)), num_ids=165) state_dict = torch.load(args.trained_model, map_location=torch.device('cpu')) modules = state_dict['model'] weights = [] for key, value in modules.items(): if not 'num_batches_tracked' in key and not 'yolo' in key: weights.append(value) print(f"{key} {value.size()}") i = 0 for name, module in net.named_modules(): if args.ebo and name is 'pair1': break if isinstance(module, torch.nn.Conv2d): print(f"write to {name}.weight") assert module.weight.size() == weights[i].size() module.weight.data = weights[i].data.clone() i += 1 if module.bias is not None: print(f"write to {name}.bias") assert module.bias.size() == weights[i].size() module.bias.data = weights[i].data.clone() i += 1 elif isinstance(module, torch.nn.BatchNorm2d): print(f"write to {name}") assert module.weight.size() == weights[i].size() module.weight.data = weights[i].data.clone() i += 1 assert module.bias.size() == weights[i].size() module.bias.data = weights[i].data.clone() i += 1 assert module.running_mean.size() == weights[i].size() module.running_mean.data = weights[i].data.clone() i += 1 assert module.running_var.size() == weights[i].size() module.running_var.data = weights[i].data.clone() i += 1 elif isinstance(module, torch.nn.Linear): print(f"write to {name}") assert module.weight.size() == weights[i].size() module.weight.data = weights[i].data.clone() i += 1 assert module.bias.size() == weights[i].size() module.bias.data = weights[i].data.clone() torch.save(net.state_dict(), args.exported_model)
def toonnx(args): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if not args.full_model: if args.backbone == 'darknet': model = darknet.DarkNet(np.random.randint(0, 100, (12, 2))).to(device) elif args.backbone == 'shufflenetv2': model = shufflenetv2.ShuffleNetV2(np.random.randint( 0, 100, (12, 2)), model_size=args.thin).to(device) else: print('unknown backbone architecture!') sys.exit(0) if args.pytorch_model: model_state_dict = model.state_dict() state_dict = torch.load(args.pytorch_model, map_location=device) state_dict = { k: v for k, v in state_dict.items() if k in model_state_dict } state_dict = collections.OrderedDict(state_dict) model_state_dict.update(state_dict) model.load_state_dict(model_state_dict) else: print('Warning: this function has not been tested yet!') model = torch.load(args.pytorch_model) model.eval() dummy_input = torch.rand(1, 3, args.insize[0], args.insize[1], device=device) onnx.export(model, dummy_input, args.onnx_model, verbose=True, input_names=['data'], output_names=['out1', 'out2', 'out3'], opset_version=11) session = ort.InferenceSession(args.onnx_model) outputs = session.run(None, {'data': dummy_input.cpu().numpy()}) for i, output in enumerate(outputs): print('branch {} output size is {}'.format(i, output.shape))
def main(args): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if args.backbone == 'darknet': model = darknet.DarkNet(np.random.randint(0, 100, (12, 2))).to(device) elif args.backbone == 'shufflenetv2': model = shufflenetv2.ShuffleNetV2(np.random.randint(0, 100, (12, 2)), model_size=args.thin).to(device) else: print('unknown backbone architecture!') sys.exit(0) # load state dict except the classifier layer model_dict = model.state_dict() trained_model_dict = torch.load(os.path.join(args.model), map_location='cpu') trained_model_dict = {k : v for (k, v) in trained_model_dict.items() if k in model_dict} trained_model_dict = collections.OrderedDict(trained_model_dict) model_dict.update(trained_model_dict) model.load_state_dict(model_dict) model.eval() if '320x576' in args.insize: anchors = ((6,16), (8,23), (11,32), (16,45), (21,64), (30,90), (43,128), (60,180), (85,255), (120,360), (170,420), (340,320)) elif '480x864' in args.insize: anchors = ((6,19), (9,27), (13,38), (18,54), (25,76), (36,107), (51,152), (71,215), (102,305), (143,429), (203,508), (407,508)) elif '608x1088' in args.insize: anchors = ((8,24), (11,34), (16,48), (23,68), (32,96), (45,135), (64,192), (90,271), (128,384), (180,540), (256,640), (512,640)) h, w = [int(s) for s in args.insize.split('x')] decoder = jde.JDEcoder((h, w), embd_dim=args.embedding) tracker = JDETracker() if os.path.isfile(args.img_path): dataloader = dataset.VideoLoader(args.img_path, (h,w,3)) else: dataloader = dataset.ImagesLoader(args.img_path, (h,w,3), formats=['*.jpg', '*.png']) strs = re.split(r'[\\, /]', args.img_path) imgpath = os.path.join(args.workspace, 'result', strs[-3], 'img') mkdir(imgpath) traj_path = os.path.join(args.workspace, 'result', '{}.txt'.format(strs[-3])) os.system('rm -f {}'.format(os.path.join(imgpath, '*'))) for i, (path, im, lb_im) in enumerate(dataloader): input = torch.from_numpy(lb_im).unsqueeze(0).to(device) with torch.no_grad(): outputs = model(input) outputs = decoder(outputs) print('{} {} {} {}'.format(path, im.shape, lb_im.shape, outputs.size()), end=' ') outputs = nonmax_suppression(outputs, args.score_thresh, args.iou_thresh)[0] if outputs is None: print('no object detected!') segments = re.split(r'[\\, /]', path) cv2.imwrite(os.path.join(imgpath, segments[-1]), im) continue print('{}'.format(outputs.size()), end=' ') outputs[:, :4] = ltrb_net2img(outputs[:, :4], (h,w), im.shape[:2]) if not args.only_detect: trajectories = tracker.update(outputs.numpy()) print('{}'.format(len(trajectories))) result = overlap_trajectory(trajectories, im) save_trajectories(traj_path, trajectories, i + 1) else: print('') result = overlap(outputs, im) segments = re.split(r'[\\, /]', path) cv2.imwrite(os.path.join(imgpath, segments[-1]), result) os.system('ffmpeg -f image2 -i {} {}.mp4 -y'.format(os.path.join(imgpath, '%06d.jpg'), os.path.join(args.workspace, 'result', strs[-3])))
'-dm', type=str, dest='dm', default='darknet.weights', help='darknet-format model file') parser.add_argument('--load-backbone-only', '-lbo', dest='lbo', help='only load the backbone', action='store_true') args = parser.parse_args() in_size = [int(insz) for insz in args.in_size.split(',')] anchors = np.loadtxt(os.path.join(args.dataset, 'anchors.txt')) model = darknet.DarkNet(anchors, in_size=in_size, num_classes=args.num_classes) with open(args.dm, 'rb') as file: major = np.fromfile(file, dtype=np.int32, count=1) minor = np.fromfile(file, dtype=np.int32, count=1) revision = np.fromfile(file, dtype=np.int32, count=1) seen = np.fromfile(file, dtype=np.int64, count=1) print( f'darknet model version:{major.data[0]}.{minor.data[0]}.{revision.data[0]}' ) last_conv = None last_name = None for name, module in model.named_modules(): if args.lbo and name is 'pair1':
import darknet import dataset if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--pytorch-model', '-pm', type=str, dest='pm', help='pytorch-format model file') parser.add_argument('--dataset', type=str, default='', help='dataset path') parser.add_argument('--num-classes', type=int, default=1, help='number of classes') parser.add_argument('--darknet-model', '-dm', type=str, dest='dm', default='darknet.weights', help='darknet-format model file') parser.add_argument('--load-backbone-only', '-lbo', dest='lbo', help='only load the backbone', action='store_true') args = parser.parse_args() dataset = dataset.CustomDataset(args.dataset, 'train') num_ids = dataset.max_id + 2 print(num_ids) model = darknet.DarkNet(np.random.randint(0, 100, (12, 2)), num_classes=args.num_classes, num_ids=num_ids) with open(args.dm, 'rb') as file: major = np.fromfile(file, dtype=np.int32, count=1) minor = np.fromfile(file, dtype=np.int32, count=1) revision = np.fromfile(file, dtype=np.int32, count=1) seen = np.fromfile(file, dtype=np.int64, count=1) print(f'darknet model version:{major.data[0]}.{minor.data[0]}.{revision.data[0]}') last_conv = None last_name = None for name, module in model.named_modules(): if args.lbo and name is 'pair1': break if isinstance(module, torch.nn.Conv2d): last_conv = module
parser.add_argument('--model', type=str, default='', help='model file') parser.add_argument('--dataset', type=str, default='', help='dataset path') parser.add_argument('--num-classes', type=int, default=3, help='number of classes') parser.add_argument('--pruned-model', '-pm', action='store_true') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if not args.pruned_model: in_size = [int(insz) for insz in args.in_size.split(',')] anchors = np.loadtxt(os.path.join(args.dataset, 'anchors.txt')) model = darknet.DarkNet(anchors, in_size, num_classes=args.num_classes).to(device) model.load_state_dict(torch.load(args.model, map_location=device)) else: model = torch.load(args.model, map_location=device) model.eval() nparams = 0 for name, module in model.named_modules(): if isinstance(module, torch.nn.Conv2d): nparams += module.weight.numel() if module.bias is not None: nparams += module.bias.numel() if isinstance(module, torch.nn.BatchNorm2d): nparams += module.weight.numel() nparams += module.bias.numel() nparams += module.running_mean.numel()
help='backbone architecture[darknet53(default),shufflenetv2]') parser.add_argument( '--thin', type=str, default='2.0x', help= 'shufflenetv2 thin, default is 2.0x, candidates are 0.5x, 1.0x, 1.5x') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') in_size = [int(insz) for insz in args.in_size.split(',')] if not args.pruned_model: dummy_anchors = np.random.randint(0, 100, (12, 2)) if args.backbone == 'darknet53': model = darknet.DarkNet(dummy_anchors).to(device) elif args.backbone == 'shufflenetv2': model = shufflenetv2.ShuffleNetV2(dummy_anchors, model_size=args.thin).to(device) else: print('unknown backbone architecture!') sys.exit(0) model_dict = model.state_dict() trained_model_dict = torch.load(args.model, map_location='cpu') trained_model_dict = { k: v for (k, v) in trained_model_dict.items() if k in model_dict } trained_model_dict = collections.OrderedDict(trained_model_dict) model_dict.update(trained_model_dict) model.load_state_dict(model_dict)
def train(args): utils.make_workspace_dirs(args.workspace) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') anchors = np.loadtxt(os.path.join(args.dataset, 'anchors.txt')) scale_sampler = utils.TrainScaleSampler(args.in_size, args.scale_step, args.rescale_freq) shared_size = torch.IntTensor(args.in_size).share_memory_() logger = utils.get_logger(path=os.path.join(args.workspace, 'log.txt')) torch.backends.cudnn.benchmark = True dataset = ds.CustomDataset(args.dataset, 'train') collate_fn = partial(ds.collate_fn, in_size=shared_size, train=True) data_loader = torch.utils.data.DataLoader(dataset, args.batch_size, True, num_workers=args.workers, collate_fn=collate_fn, pin_memory=args.pin, drop_last=True) num_ids = dataset.max_id + 2 if args.backbone == 'darknet': model = darknet.DarkNet(anchors, num_classes=args.num_classes, num_ids=num_ids).to(device) elif args.backbone == 'shufflenetv2': model = shufflenetv2.ShuffleNetV2(anchors, num_classes=args.num_classes, num_ids=num_ids, model_size=args.thin).to(device) else: print('unknown backbone architecture!') sys.exit(0) if args.checkpoint: model.load_state_dict(torch.load(args.checkpoint)) params = [p for p in model.parameters() if p.requires_grad] if args.optim == 'sgd': optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) else: optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay) if args.freeze_bn: for name, param in model.named_parameters(): if 'norm' in name: param.requires_grad = False logger.info('freeze {}'.format(name)) else: param.requires_grad = True trainer = f'{args.workspace}/checkpoint/trainer-ckpt.pth' if args.resume: trainer_state = torch.load(trainer) optimizer.load_state_dict(trainer_state['optimizer']) if -1 in args.milestones: args.milestones = [int(args.epochs * 0.5), int(args.epochs * 0.75)] lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=args.milestones, gamma=args.lr_gamma) start_epoch = 0 if args.resume: start_epoch = trainer_state['epoch'] + 1 lr_scheduler.load_state_dict(trainer_state['lr_scheduler']) logger.info(args) logger.info('Start training from epoch {}'.format(start_epoch)) model_path = f'{args.workspace}/checkpoint/{args.savename}-ckpt-%03d.pth' size = shared_size.numpy().tolist() for epoch in range(start_epoch, args.epochs): model.train() logger.info(('%8s%10s%10s' + '%10s' * 8) % ('Epoch', 'Batch', 'SIZE', 'LBOX', 'LCLS', 'LIDE', 'LOSS', 'SB', 'SC', 'SI', 'LR')) rmetrics = defaultdict(float) optimizer.zero_grad() for batch, (images, targets) in enumerate(data_loader): warmup = min(args.warmup, len(data_loader)) if epoch == 0 and batch <= warmup: lr = args.lr * (batch / warmup)**4 for g in optimizer.param_groups: g['lr'] = lr loss, metrics = model(images.to(device), targets.to(device), size) loss.backward() if args.sparsity: model.correct_bn_grad(args.lamb) num_batches = epoch * len(data_loader) + batch + 1 if ((batch + 1) % args.accumulated_batches == 0) or (batch == len(data_loader) - 1): optimizer.step() optimizer.zero_grad() for k, v in metrics.items(): rmetrics[k] = (rmetrics[k] * batch + metrics[k]) / (batch + 1) fmt = tuple([('%g/%g') % (epoch, args.epochs), ('%g/%g') % (batch, len(data_loader)), ('%gx%g') % (size[0], size[1])] + \ list(rmetrics.values()) + [optimizer.param_groups[0]['lr']]) if batch % args.print_interval == 0: logger.info(('%8s%10s%10s' + '%10.3g' * (len(rmetrics.values()) + 1)) % fmt) size = scale_sampler(num_batches) shared_size[0], shared_size[1] = size[0], size[1] torch.save(model.state_dict(), f"{model_path}" % epoch) torch.save( { 'epoch': epoch, 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict() }, trainer) if epoch >= args.eval_epoch: pass lr_scheduler.step()
import os import torch import argparse import yolov3 import darknet if __name__ == '__main__': parser = argparse.ArgumentParser( description='compare the results of reference and own implementation') parser.add_argument('--path', type=str, help='path to the reference data') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = darknet.DarkNet().to(device) model.load_state_dict( torch.load(os.path.join(args.path, 'jde.pth'), map_location='cpu')) model.eval() x = torch.load(os.path.join(args.path, 'input.pt'), map_location=device) with torch.no_grad(): outputs = model(x) refers = [] refers.append( torch.load(os.path.join(args.path, 'out1.pt'), map_location=device)) refers.append( torch.load(os.path.join(args.path, 'out2.pt'), map_location=device)) refers.append( torch.load(os.path.join(args.path, 'out3.pt'), map_location=device))
print(args) in_size = [int(size) for size in args.in_size.split(',')] if os.path.isfile(args.cmodel): file = open(args.cmodel, 'rb') net = pb.NetParameter() net.ParseFromString(file.read()) file.close() prototxt = args.cmodel.replace('caffemodel', 'prototxt') write_prototxt(net, prototxt, in_size) sys.exit(0) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if not args.pruned_model: dummy_anchors = np.random.randint(0, 100, (12, 2)) model = darknet.DarkNet(dummy_anchors) # load state dict except the classifier layer model_dict = model.state_dict() trained_model_dict = torch.load(args.pmodel, map_location='cpu') trained_model_dict = { k: v for (k, v) in trained_model_dict.items() if k in model_dict } trained_model_dict = collections.OrderedDict(trained_model_dict) model_dict.update(trained_model_dict) model.load_state_dict(model_dict) else: model = torch.load(args.pmodel, map_location=device) # we have to process route layer specially! route_bottoms = [['cbrl7_relu'], ['upsample1', 'stage4_7_ew'],
help='backbone architecture[darknet53(default),shufflenetv2]') parser.add_argument('--pruned-model', action='store_true') args = parser.parse_args() print(args) in_size = [int(insz) for insz in args.in_size.split(',')] classnames = utils.load_class_names( os.path.join(args.dataset, 'classes.txt')) anchors = np.loadtxt(os.path.join(args.dataset, 'anchors.txt')) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') norm = 255 if args.backbone == 'darknet53' else 1 if not args.pruned_model: if args.backbone == 'darknet53': model = darknet.DarkNet(anchors, in_size=in_size, num_classes=len(classnames)).to(device) elif args.backbone == 'shufflenetv2': model = shufflenetv2.ShuffleNetV2( anchors, in_size=in_size, num_classes=len(classnames)).to(device) else: print('unknown backbone architecture!') sys.exit(0) model.load_state_dict(torch.load(args.model, map_location=device)) else: model = torch.load(args.model, map_location=device) model.eval() decoder = yolov3.YOLOv3EvalDecoder(in_size, len(classnames), anchors) transform = dataset.get_transform(train=False,
def train(args): utils.make_workspace_dirs(args.workspace) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') anchors = np.loadtxt(args.anchors) if args.anchors else None scale_sampler = utils.TrainScaleSampler(args.in_size, args.scale_step, args.rescale_freq) shared_size = torch.IntTensor(args.in_size).share_memory_() logger = utils.get_logger(path=os.path.join(args.workspace, 'log.txt')) torch.backends.cudnn.benchmark = True dataset = ds.HotchpotchDataset(args.dataset_root, './data/train.txt', args.backbone) collate_fn = partial(ds.collate_fn, in_size=shared_size, train=True) data_loader = torch.utils.data.DataLoader(dataset, args.batch_size, True, num_workers=args.workers, collate_fn=collate_fn, pin_memory=args.pin, drop_last=True) num_ids = int(dataset.max_id + 1) if args.backbone == 'darknet': model = darknet.DarkNet(anchors, num_classes=args.num_classes, num_ids=num_ids).to(device) elif args.backbone == 'shufflenetv2': model = shufflenetv2.ShuffleNetV2(anchors, num_classes=args.num_classes, num_ids=num_ids, model_size=args.thin, box_loss=args.box_loss, cls_loss=args.cls_loss).to(device) else: print('unknown backbone architecture!') sys.exit(0) if args.checkpoint: model.load_state_dict(torch.load(args.checkpoint)) lr_min = 0.00025 params = [p for p in model.parameters() if p.requires_grad] backbone_neck_params, detection_params, identity_params = grouping_model_params( model) if args.optim == 'sgd': # optimizer = torch.optim.SGD(params, lr=args.lr, # momentum=args.momentum, weight_decay=args.weight_decay) optimizer = torch.optim.SGD([{ 'params': backbone_neck_params }, { 'params': detection_params, 'lr': args.lr * args.lr_coeff[1] }, { 'params': identity_params, 'lr': args.lr * args.lr_coeff[2] }], lr=(args.lr - lr_min), momentum=args.momentum, weight_decay=args.weight_decay) else: optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay) if args.freeze_bn: for name, param in model.named_parameters(): if 'norm' in name: param.requires_grad = False logger.info('freeze {}'.format(name)) else: param.requires_grad = True trainer = f'{args.workspace}/checkpoint/trainer-ckpt.pth' if args.resume: trainer_state = torch.load(trainer) optimizer.load_state_dict(trainer_state['optimizer']) def lr_lambda(batch): return 0.5 * math.cos( (batch % len(data_loader)) / (len(data_loader) - 1) * math.pi) lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) start_epoch = 0 logger.info(args) logger.info('Start training from epoch {}'.format(start_epoch)) model_path = f'{args.workspace}/checkpoint/{args.savename}-ckpt-%03d.pth' size = shared_size.numpy().tolist() for epoch in range(start_epoch, args.epochs): model.train() logger.info(('%8s%10s%10s' + '%10s' * 8) % ('Epoch', 'Batch', 'SIZE', 'LBOX', 'LCLS', 'LIDE', 'LOSS', 'SBOX', 'SCLS', 'SIDE', 'LR')) rmetrics = defaultdict(float) optimizer.zero_grad() for batch, (images, targets) in enumerate(data_loader): for i, g in enumerate(optimizer.param_groups): g['lr'] += (args.lr - lr_min) * 0.5 + lr_min loss, metrics = model(images.to(device), targets.to(device), size) loss.backward() if args.sparsity: model.correct_bn_grad(args.lamb) num_batches = epoch * len(data_loader) + batch + 1 if ((batch + 1) % args.accumulated_batches == 0) or (batch == len(data_loader) - 1): optimizer.step() optimizer.zero_grad() for k, v in metrics.items(): rmetrics[k] = (rmetrics[k] * batch + metrics[k]) / (batch + 1) fmt = tuple([('%g/%g') % (epoch, args.epochs), ('%g/%g') % (batch, len(data_loader)), ('%gx%g') % (size[0], size[1])] + \ list(rmetrics.values()) + [optimizer.param_groups[0]['lr']]) if batch % args.print_interval == 0: logger.info(('%8s%10s%10s' + '%10.3g' * (len(rmetrics.values()) + 1)) % fmt) size = scale_sampler(num_batches) shared_size[0], shared_size[1] = size[0], size[1] lr_scheduler.step() torch.save(model.state_dict(), f"{model_path}" % epoch) torch.save( { 'epoch': epoch, 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict() }, trainer) if epoch >= args.eval_epoch: pass
help='input size') parser.add_argument('--dataset', type=str, default='', help='dataset path') parser.add_argument('--onnxname', type=str, default='model.onnx', help='exported ONNX filename') parser.add_argument('--pruned-model', action='store_true', help='the model has been pruned') args = parser.parse_args() insize = [int(s) for s in args.insize.split(',')] device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if not args.pruned_model: anchors = np.loadtxt(os.path.join(args.dataset, 'anchors.txt')) model = net.DarkNet(anchors, insize, args.num_classes).to(device) model.load_state_dict(torch.load(args.model, map_location=device)) else: model = torch.load(args.model, map_location=device) input_names = ["data"] dummy_input = torch.randn(1, 3, insize[0], insize[1], device=device) torch.onnx.export(model, dummy_input, args.onnxname, verbose=True, input_names=input_names) # filename, extension = splitext(args.onnxname) # os.system(f'python -m onnxsim {args.onnxname} {args.onnxname}') # os.system(f'./model/onnx/onnx2ncnn {args.onnxname} {filename}.param {filename}.bin')
def main(args): try: mp.set_start_method('spawn') except RuntimeError: pass utils.make_workspace_dirs(args.workspace) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') in_size = [int(s) for s in args.in_size.split(',')] scale_step = [int(ss) for ss in args.scale_step.split(',')] anchors = np.loadtxt(os.path.join(args.dataset, 'anchors.txt')) scale_sampler = utils.TrainScaleSampler(scale_step, args.rescale_freq) shared_size = torch.IntTensor(in_size).share_memory_() dataset = ds.CustomDataset(args.dataset, 'train') collate_fn = partial(ds.collate_fn, in_size=shared_size, train=True) data_loader = torch.utils.data.DataLoader(dataset, args.batch_size, True, num_workers=args.workers, collate_fn=collate_fn, pin_memory=args.pin) model = darknet.DarkNet(anchors, in_size, num_classes=args.num_classes).to(device) if args.checkpoint: print(f'load {args.checkpoint}') model.load_state_dict(torch.load(args.checkpoint)) if args.sparsity: model.load_prune_permit('model/prune_permit.json') criterion = yolov3.YOLOv3Loss(args.num_classes, anchors) decoder = yolov3.YOLOv3EvalDecoder(in_size, args.num_classes, anchors) params = [p for p in model.parameters() if p.requires_grad] if args.optim == 'sgd': optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) else: optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay) trainer = f'{args.workspace}/checkpoint/trainer-ckpt.pth' if args.resume: trainer_state = torch.load(trainer) optimizer.load_state_dict(trainer_state['optimizer']) milestones = [int(ms) for ms in args.milestones.split(',')] def lr_lambda(iter): if iter < args.warmup: return pow(iter / args.warmup, 4) factor = 1 for i in milestones: factor *= pow(args.lr_gamma, int(iter > i)) return factor lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) if args.resume: start_epoch = trainer_state['epoch'] + 1 lr_scheduler.load_state_dict(trainer_state['lr_scheduler']) else: start_epoch = 0 print(f'Start training from epoch {start_epoch}') for epoch in range(start_epoch, args.epochs): msgs = train_one_epoch(model, criterion, optimizer, lr_scheduler, data_loader, epoch, args.interval, shared_size, scale_sampler, device, args.sparsity, args.lamb) utils.print_training_message(args.workspace, epoch + 1, msgs, args.batch_size) torch.save( model.state_dict(), f"{args.workspace}/checkpoint/{args.savename}-ckpt-%03d.pth" % epoch) torch.save( { 'epoch': epoch, 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict() }, trainer) if epoch >= args.eval_epoch: pass