from collections import OrderedDict import dataset if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--src-model', '-s', type=str, help='source model file') parser.add_argument('--dst-model', '-d', type=str, help='destination model file') parser.add_argument('--num-classes', type=int, default=1, help='number of classes') parser.add_argument('--dataset', type=str, default='', help='dataset path') args = parser.parse_args() print(args) # dataset = dataset.CustomDataset(args.dataset, 'train') # num_ids = dataset.max_id + 2 dataset = dataset.HotchpotchDataset('/data/tseng/dataset/jde', './data/train.txt') num_ids = int(dataset.max_id + 1) 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)
import sys import torch import pickle import numpy as np from functools import partial import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D sys.path.append('.') import dataset as ds boxw0 = [] boxh0 = [] cachefile = './wh.pkl' if not os.path.isfile(cachefile): dataset = ds.HotchpotchDataset('/data/tseng/dataset/jde', './data/train.txt', 'shufflenetv2', False) in_size = torch.IntTensor([320, 576]) collate_fn = partial(ds.collate_fn, in_size=in_size, train=False) data_loader = torch.utils.data.DataLoader(dataset, collate_fn=collate_fn) for batch, (images, targets) in enumerate(data_loader): n, c, h, w = images.size() if targets.size(0) > 0: tw = (targets[:, 5] * w).numpy().round().astype(np.int) th = (targets[:, 6] * h).numpy().round().astype(np.int) mask = (tw < 8) | (th < 8) if mask.sum() > 0: with open('error_box.txt', 'a') as file: file.write('{}\n'.format(targets.numpy())) file.write('{}x{}\n'.format(h, w)) file.write('{}\n'.format(tw)) file.write('{}\n\n\n'.format(th))
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
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.HotchpotchDataset(args.dataset, './data/train.txt') num_ids = int(dataset.max_id + 1) 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]}' )
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', args.backbone) dataset = ds.HotchpotchDataset('/data/tseng/dataset/jde', './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 = dataset.max_id + 2 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).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()