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)