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
