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 shufflenetv2(width_mult=1., pretrained=True): model = ShuffleNetV2.ShuffleNetV2(input_size=160, width_mult=width_mult) if pretrained: if width_mult == 1.: model_file = osp.join(this_dir, 'model/97_160_2.pth') elif width_mult == .5: model_file = osp.join( this_dir, '../models/ShuffleNetV2/shufflenetv2_x0.5_60.646_81.696.pth.tar' ) else: raise ValueError( 'width_mult = {} is not support pretrained model'.format( width_mult)) assert osp.exists(model_file), '{} is not exists!'.format(model_file) print('shufflenetv2 load state dict: ', model_file) state = torch.load(model_file) model.load_state_dict(state['net'], strict=False) return model
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])))
dataset = dataset.CustomDataset(args.dataset, 'train') num_ids = dataset.max_id + 2 print(num_ids) if '0.5x' in args.src_model: model_size = '0.5x' elif '1.0x' in args.src_model: model_size = '1.0x' elif '1.5x' in args.src_model: model_size = '1.5x' elif '2.0x' in args.src_model: model_size = '2.0x' anchors = np.random.randint(low=10, high=150, size=(12, 2)) model = shufflenetv2.ShuffleNetV2(anchors, num_classes=args.num_classes, num_ids=num_ids, model_size=model_size) checkpoint = torch.load(args.src_model, map_location='cpu') if 'state_dict' in checkpoint: checkpoint = dict(checkpoint['state_dict']) src_param = list() for k, v in checkpoint.items(): if 'conv_last' in k: break if not 'num_batches_tracked' in k: src_param.append(v) src_layer_index = 0 for name, module in model.named_modules(): if 'stage5' in name:
parser.add_argument('--pruned-model', '-pm', action='store_true') parser.add_argument( '--backbone', type=str, default='darknet53', help='backbone architecture[darknet53(default),shufflenetv2]') 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: anchors = np.loadtxt(os.path.join(args.dataset, 'anchors.txt')) if args.backbone == 'darknet53': model = darknet.DarkNet(anchors, in_size, num_classes=args.num_classes).to(device) elif args.backbone == 'shufflenetv2': model = shufflenetv2.ShuffleNetV2( anchors, in_size=in_size, num_classes=args.num_classes).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) input = torch.randn(1, 3, in_size[0], in_size[1]).to(device) flops, params = profile(model, inputs=(input, )) print(f'flops={flops}, params={params}')
help='nms iou threshold, default=0.4, it must be in [0,1]') parser.add_argument('--save-path', type=str, default=os.path.join(os.getcwd(), 'mining'), help='path to the mining samples') args = parser.parse_args() return args def mkdir(path): if not os.path.exists(path): os.makedirs(path) if __name__ == '__main__': args = parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = shufflenetv2.ShuffleNetV2().to(device) 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) model.eval() h, w = args.in_size decoder = jde.JDEcoder(args.in_size, embd_dim=args.embedding) dataloader = dataset.VideoLoader(args.video, (h, w, 3)) mkdir(os.path.join(args.save_path, 'images'))
default=os.path.join(os.getcwd(), 'mining'), help='path to the mining samples') args = parser.parse_args() return args def mkdir(path): if not os.path.exists(path): os.makedirs(path) if __name__ == '__main__': args = parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = shufflenetv2.ShuffleNetV2(np.random.randint(0, 100, (12, 2)), model_size='1.0x').to(device) 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) model.eval() h, w = args.in_size decoder = jde.JDEcoder(args.in_size, embd_dim=args.embedding)
type=str, default='shufflenetv2', help='backbone architecture, default is shufflenetv2' 'other option is darknet') parser.add_argument( '--thin', type=str, default='0.5x', help='shufflenetv2 backbone thin, default is 0.5x, other options' 'are 1.0x, 1.5x, and 2.0x') args = parser.parse_args() # construct JDE model anchors = np.random.randint(low=0, high=200, size=(12, 2)) if args.backbone == 'shufflenetv2': model = shufflenetv2.ShuffleNetV2(anchors, model_size=args.thin) elif args.backbone == 'darknet': raise NotImplementedError else: raise NotImplementedError # load weights if PyTorch model was given if args.pytorch_model: state_dict = model.state_dict() train_state_dict = torch.load(args.pytorch_model) # remove identifier classifier from train_state_dict train_state_dict = { k: v for k, v in train_state_dict.items() if k in state_dict.keys() } train_state_dict = collections.OrderedDict(train_state_dict)
parser.add_argument('--model', type=str, help='path to the model') parser.add_argument('--input', type=str, help='path to the input tensor') parser.add_argument('--output', type=str, help='path to the tensorrt inference output tensor') parser.add_argument('--index', type=int, default=0, help='the output index') args = parser.parse_args() 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)) model = shufflenetv2.ShuffleNetV2(anchors, model_size='0.5x') state_dict = model.state_dict() train_state_dict = torch.load(args.model) train_state_dict = { k: v for k, v in train_state_dict.items() if k in state_dict.keys() } train_state_dict = collections.OrderedDict(train_state_dict) state_dict.update(train_state_dict) model.load_state_dict(state_dict) model.eval() x = torch.from_numpy(np.fromfile(args.input, np.float32)).view(1, 3, 320, 576)
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 sys import glob import torch import argparse import collections from torch import nn from collections import defaultdict sys.path.append('.') import shufflenetv2 parser = argparse.ArgumentParser() parser.add_argument('--model-path', '-mp', type=str, help='path to models') args = parser.parse_args() avr_model = shufflenetv2.ShuffleNetV2(model_size='1.0x') for name, module in avr_model.named_modules(): if isinstance(module, nn.Conv2d): module.weight.data[...] = 0 if module.bias is not None: module.bias.data[...] = 0 elif isinstance(module, nn.BatchNorm2d): module.weight.data[...] = 0 module.bias.data[...] = 0 module.running_mean.data[...] = 0 module.running_var.data[...] = 0 paths = sorted(glob.glob(os.path.join(args.model_path, '*.pth'))) paths = [path for path in paths if 'trainer-ckpt' not in path] print(paths) params = defaultdict(list)
import os.path as osp sys.path.append('.') import jde import shufflenetv2 parser = argparse.ArgumentParser() parser.add_argument('--weight', '-w', type=str, help='weight path') parser.add_argument('--data', '-d', type=str, help='path to tensorrt generated data') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = shufflenetv2.ShuffleNetV2(model_size='1.0x').to(device) model_dict = model.state_dict() trained_model_dict = torch.load(args.weight, 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() decoder = jde.JDEcoder((320, 576))
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, net_w=in_size[0], net_h=in_size[1]) FloatTensor = torch.cuda.FloatTensor if torch.cuda.is_available( ) else torch.FloatTensor
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
'--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) else: model = torch.load(args.model, map_location=device)