SHARE
TWEET

Untitled

a guest Jan 7th, 2019 99 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. """Train YOLOv3 with random shapes."""
  2. import argparse
  3. import os
  4. import logging
  5. import time
  6. import warnings
  7. import numpy as np
  8. import mxnet as mx
  9. from mxnet import nd
  10. from mxnet import gluon
  11. from mxnet import autograd
  12. import gluoncv as gcv
  13. from gluoncv import data as gdata
  14. from gluoncv import utils as gutils
  15. from gluoncv.model_zoo import get_model
  16. from gluoncv.data.batchify import Tuple, Stack, Pad
  17. from gluoncv.data.transforms.presets.yolo import YOLO3DefaultTrainTransform
  18. from gluoncv.data.transforms.presets.yolo import YOLO3DefaultValTransform
  19. from gluoncv.data.dataloader import RandomTransformDataLoader
  20. from gluoncv.utils.metrics.voc_detection import VOC07MApMetric
  21. from gluoncv.utils.metrics.coco_detection import COCODetectionMetric
  22. from gluoncv.utils import LRScheduler
  23.  
  24. from gluoncv.utils import download, viz
  25. from matplotlib import pyplot as plt
  26.    
  27. def parse_args():
  28.     parser = argparse.ArgumentParser(description='Train YOLO networks with random input shape.')
  29.     parser.add_argument('--network', type=str, default='yolo3_darknet53_coco',
  30.                         help="Base network name which serves as feature extraction base.")
  31.     parser.add_argument('--data-shape', type=int, default=320,
  32.                         help="Input data shape for evaluation, use 320, 416, 608... " +
  33.                              "Training is with random shapes from (320 to 608).")
  34.     parser.add_argument('--batch-size', type=int, default=4, help='Training mini-batch size')
  35.     parser.add_argument('--num-workers', '-j', dest='num_workers', type=int,
  36.                         default=0, help='Number of data workers, you can use larger '
  37.                         'number to accelerate data loading, if you CPU and GPUs are powerful.')
  38.     parser.add_argument('--gpus', type=str, default='0',
  39.                         help='Training with GPUs, you can specify 1,3 for example.')
  40.     parser.add_argument('--epochs', type=int, default=1,
  41.                         help='Training epochs.')
  42.     parser.add_argument('--resume', type=str, default='',
  43.                         help='Resume from previously saved parameters if not None. '
  44.                         'For example, you can resume from ./yolo3_xxx_0123.params')
  45.     parser.add_argument('--start-epoch', type=int, default=0,
  46.                         help='Starting epoch for resuming, default is 0 for new training.'
  47.                         'You can specify it to 100 for example to start from 100 epoch.')
  48.     parser.add_argument('--lr', type=float, default=0.001,
  49.                         help='Learning rate, default is 0.001')
  50.     parser.add_argument('--lr-mode', type=str, default='step',
  51.                         help='learning rate scheduler mode. options are step, poly and cosine.')
  52.     parser.add_argument('--lr-decay', type=float, default=0.1,
  53.                         help='decay rate of learning rate. default is 0.1.')
  54.     parser.add_argument('--lr-decay-period', type=int, default=0,
  55.                         help='interval for periodic learning rate decays. default is 0 to disable.')
  56.     parser.add_argument('--lr-decay-epoch', type=str, default='160,180',
  57.                         help='epochs at which learning rate decays. default is 160,180.')
  58.     parser.add_argument('--warmup-lr', type=float, default=0.0,
  59.                         help='starting warmup learning rate. default is 0.0.')
  60.     parser.add_argument('--warmup-epochs', type=int, default=0,
  61.                         help='number of warmup epochs.')
  62.     parser.add_argument('--momentum', type=float, default=0.9,
  63.                         help='SGD momentum, default is 0.9')
  64.     parser.add_argument('--wd', type=float, default=0.0005,
  65.                         help='Weight decay, default is 5e-4')
  66.     parser.add_argument('--log-interval', type=int, default=100,
  67.                         help='Logging mini-batch interval. Default is 100.')
  68.     parser.add_argument('--save-prefix', type=str, default='',
  69.                         help='Saving parameter prefix')
  70.     parser.add_argument('--save-interval', type=int, default=10,
  71.                         help='Saving parameters epoch interval, best model will always be saved.')
  72.     parser.add_argument('--val-interval', type=int, default=1,
  73.                         help='Epoch interval for validation, increase the number will reduce the '
  74.                              'training time if validation is slow.')
  75.     parser.add_argument('--seed', type=int, default=233,
  76.                         help='Random seed to be fixed.')
  77.     parser.add_argument('--num-samples', type=int, default=-1,
  78.                         help='Training images. Use -1 to automatically get the number.')
  79.     parser.add_argument('--syncbn', action='store_true',
  80.                         help='Use synchronize BN across devices.')
  81.     parser.add_argument('--no-random-shape', action='store_true',
  82.                         help='Use fixed size(data-shape) throughout the training, which will be faster '
  83.                         'and require less memory. However, final model will be slightly worse.')
  84.     parser.add_argument('--no-wd', action='store_true',
  85.                         help='whether to remove weight decay on bias, and beta/gamma for batchnorm layers.')
  86.     parser.add_argument('--mixup', action='store_true',
  87.                         help='whether to enable mixup.')
  88.     parser.add_argument('--no-mixup-epochs', type=int, default=20,
  89.                         help='Disable mixup training if enabled in the last N epochs.')
  90.     parser.add_argument('--train_dataset', type=str,
  91.                         help='Location of training dataset, must be rec format')
  92.     parser.add_argument('--validate_dataset', type=str,
  93.                         help='Location of validate dataset, must be rec format')
  94.     parser.add_argument('--classes_list', type=str,
  95.                         help='Location of the classes list, every line present one class')
  96.     parser.add_argument('--pretrained', type=int, default=0,
  97.                         help='0 will use pretrained weights, other will not')
  98.    
  99.     parser.add_argument('--label-smooth', action='store_true', help='Use label smoothing.')
  100.     args = parser.parse_args()
  101.     return args
  102.    
  103. def read_classes(args):
  104.     with open(args.classes_list) as f:
  105.         return f.readlines()   
  106.  
  107. def get_dataset(args):
  108.     train_dataset = gcv.data.RecordFileDetection(args.train_dataset)
  109.     val_dataset = gcv.data.RecordFileDetection(args.validate_dataset)
  110.     classes = read_classes(args)
  111.     val_metric = VOC07MApMetric(iou_thresh=0.5, class_names=classes)
  112.    
  113.     if args.num_samples < 0:
  114.         args.num_samples = len(train_dataset)
  115.     if args.mixup:
  116.         from gluoncv.data import MixupDetection
  117.         train_dataset = MixupDetection(train_dataset)
  118.     return train_dataset, val_dataset, val_metric
  119.  
  120. def get_dataloader(net, train_dataset, val_dataset, data_shape, batch_size, num_workers, args):
  121.     """Get dataloader."""
  122.     width, height = data_shape, data_shape
  123.     batchify_fn = Tuple(*([Stack() for _ in range(6)] + [Pad(axis=0, pad_val=-1) for _ in range(1)]))  # stack image, all targets generated
  124.     if args.no_random_shape:
  125.         print("no random shape")
  126.         train_loader = gluon.data.DataLoader(
  127.             train_dataset.transform(YOLO3DefaultTrainTransform(width, height, net, mixup=args.mixup)),
  128.             batch_size, True, batchify_fn=batchify_fn, last_batch='rollover', num_workers=num_workers)
  129.     else:
  130.         print("with random shape")
  131.         transform_fns = [YOLO3DefaultTrainTransform(x * 32, x * 32, net, mixup=args.mixup) for x in range(10, 20)]
  132.         train_loader = RandomTransformDataLoader(
  133.             transform_fns, train_dataset, batch_size=batch_size, interval=10, last_batch='rollover',
  134.             shuffle=True, batchify_fn=batchify_fn, num_workers=num_workers)
  135.     val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))    
  136.     val_loader = gluon.data.DataLoader(
  137.         val_dataset.transform(YOLO3DefaultValTransform(width, height)),
  138.         batch_size, True, batchify_fn=val_batchify_fn, last_batch='keep', num_workers=num_workers)
  139.     return train_loader, val_loader
  140.  
  141. def save_params(net, best_map, current_map, epoch, save_interval, prefix):
  142.     current_map = float(current_map)
  143.     if current_map > best_map[0]:
  144.         best_map[0] = current_map
  145.         net.save_parameters('{:s}_best.params'.format(prefix, epoch, current_map))
  146.         with open(prefix+'_best_map.log', 'a') as f:
  147.             f.write('{:04d}:\t{:.4f}\n'.format(epoch, current_map))
  148.     if save_interval and epoch % save_interval == 0:
  149.         net.save_parameters('{:s}_{:04d}_{:.4f}.params'.format(prefix, epoch, current_map))
  150.  
  151. def validate(net, val_data, ctx, eval_metric):
  152.     """Test on validation dataset."""
  153.     eval_metric.reset()
  154.     # set nms threshold and topk constraint
  155.     net.set_nms(nms_thresh=0.45, nms_topk=400)
  156.     mx.nd.waitall()
  157.     net.hybridize()
  158.     for batch in val_data:
  159.         data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0, even_split=False)
  160.         label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0, even_split=False)
  161.         det_bboxes = []
  162.         det_ids = []
  163.         det_scores = []
  164.         gt_bboxes = []
  165.         gt_ids = []
  166.         gt_difficults = []
  167.         for x, y in zip(data, label):
  168.             # get prediction results
  169.             ids, scores, bboxes = net(x)
  170.             det_ids.append(ids)
  171.             det_scores.append(scores)
  172.             # clip to image size
  173.             det_bboxes.append(bboxes.clip(0, batch[0].shape[2]))
  174.             # split ground truths
  175.             gt_ids.append(y.slice_axis(axis=-1, begin=4, end=5))
  176.             gt_bboxes.append(y.slice_axis(axis=-1, begin=0, end=4))
  177.             gt_difficults.append(y.slice_axis(axis=-1, begin=5, end=6) if y.shape[-1] > 5 else None)
  178.  
  179.         # update metric
  180.         eval_metric.update(det_bboxes, det_ids, det_scores, gt_bboxes, gt_ids, gt_difficults)
  181.     return eval_metric.get()
  182.  
  183. def train(net, train_data, val_data, eval_metric, ctx, args):
  184.     """Training pipeline"""
  185.     net.collect_params().reset_ctx(ctx)
  186.     if args.no_wd:
  187.         for k, v in net.collect_params('.*beta|.*gamma|.*bias').items():
  188.             v.wd_mult = 0.0
  189.  
  190.     if args.label_smooth:
  191.         net._target_generator._label_smooth = True
  192.  
  193.     if args.lr_decay_period > 0:
  194.         lr_decay_epoch = list(range(args.lr_decay_period, args.epochs, args.lr_decay_period))
  195.     else:
  196.         lr_decay_epoch = [int(i) for i in args.lr_decay_epoch.split(',')]
  197.     lr_scheduler = LRScheduler(mode=args.lr_mode,
  198.                                baselr=args.lr,
  199.                                niters=args.num_samples // args.batch_size,
  200.                                nepochs=args.epochs,
  201.                                step=lr_decay_epoch,
  202.                                step_factor=args.lr_decay, power=2,
  203.                                warmup_epochs=args.warmup_epochs)
  204.  
  205.     trainer = gluon.Trainer(
  206.         net.collect_params(), 'sgd',
  207.         {'wd': args.wd, 'momentum': args.momentum, 'lr_scheduler': lr_scheduler},
  208.         kvstore='local')
  209.  
  210.     # targets
  211.     sigmoid_ce = gluon.loss.SigmoidBinaryCrossEntropyLoss(from_sigmoid=False)
  212.     l1_loss = gluon.loss.L1Loss()
  213.  
  214.     # metrics
  215.     obj_metrics = mx.metric.Loss('ObjLoss')
  216.     center_metrics = mx.metric.Loss('BoxCenterLoss')
  217.     scale_metrics = mx.metric.Loss('BoxScaleLoss')
  218.     cls_metrics = mx.metric.Loss('ClassLoss')
  219.  
  220.     # set up logger
  221.     logging.basicConfig()
  222.     logger = logging.getLogger()
  223.     logger.setLevel(logging.INFO)
  224.     log_file_path = args.save_prefix + '_train.log'
  225.     log_dir = os.path.dirname(log_file_path)
  226.     if log_dir and not os.path.exists(log_dir):
  227.         os.makedirs(log_dir)
  228.     fh = logging.FileHandler(log_file_path)
  229.     logger.addHandler(fh)
  230.     logger.info(args)
  231.     logger.info('Start training from [Epoch {}]'.format(args.start_epoch))
  232.     best_map = [0]
  233.     for epoch in range(args.start_epoch, args.epochs):
  234.         if args.mixup:
  235.             # TODO(zhreshold): more elegant way to control mixup during runtime
  236.             try:
  237.                 train_data._dataset.set_mixup(np.random.beta, 1.5, 1.5)
  238.             except AttributeError:
  239.                 train_data._dataset._data.set_mixup(np.random.beta, 1.5, 1.5)
  240.             if epoch >= args.epochs - args.no_mixup_epochs:
  241.                 try:
  242.                     train_data._dataset.set_mixup(None)
  243.                 except AttributeError:
  244.                     train_data._dataset._data.set_mixup(None)
  245.  
  246.         tic = time.time()
  247.         btic = time.time()
  248.         mx.nd.waitall()
  249.         net.hybridize()
  250.         for i, batch in enumerate(train_data):
  251.             #print("training batch\n", batch)
  252.             #return 0
  253.             batch_size = batch[0].shape[0]
  254.             data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0)
  255.             # objectness, center_targets, scale_targets, weights, class_targets
  256.             fixed_targets = [gluon.utils.split_and_load(batch[it], ctx_list=ctx, batch_axis=0) for it in range(1, 6)]
  257.             gt_boxes = gluon.utils.split_and_load(batch[6], ctx_list=ctx, batch_axis=0)
  258.             sum_losses = []
  259.             obj_losses = []
  260.             center_losses = []
  261.             scale_losses = []
  262.             cls_losses = []
  263.             with autograd.record():
  264.                 for ix, x in enumerate(data):
  265.                     obj_loss, center_loss, scale_loss, cls_loss = net(x, gt_boxes[ix], *[ft[ix] for ft in fixed_targets])
  266.                     sum_losses.append(obj_loss + center_loss + scale_loss + cls_loss)
  267.                     obj_losses.append(obj_loss)
  268.                     center_losses.append(center_loss)
  269.                     scale_losses.append(scale_loss)
  270.                     cls_losses.append(cls_loss)
  271.                 autograd.backward(sum_losses)
  272.             lr_scheduler.update(i, epoch)
  273.             trainer.step(batch_size)
  274.             obj_metrics.update(0, obj_losses)
  275.             center_metrics.update(0, center_losses)
  276.             scale_metrics.update(0, scale_losses)
  277.             cls_metrics.update(0, cls_losses)
  278.             if args.log_interval and not (i + 1) % args.log_interval:
  279.                 name1, loss1 = obj_metrics.get()
  280.                 name2, loss2 = center_metrics.get()
  281.                 name3, loss3 = scale_metrics.get()
  282.                 name4, loss4 = cls_metrics.get()
  283.                 logger.info('[Epoch {}][Batch {}], LR: {:.2E}, Speed: {:.3f} samples/sec, {}={:.3f}, {}={:.3f}, {}={:.3f}, {}={:.3f}'.format(
  284.                     epoch, i, trainer.learning_rate, batch_size/(time.time()-btic), name1, loss1, name2, loss2, name3, loss3, name4, loss4))
  285.             btic = time.time()
  286.  
  287.         name1, loss1 = obj_metrics.get()
  288.         name2, loss2 = center_metrics.get()
  289.         name3, loss3 = scale_metrics.get()
  290.         name4, loss4 = cls_metrics.get()
  291.         logger.info('[Epoch {}] Training cost: {:.3f}, {}={:.3f}, {}={:.3f}, {}={:.3f}, {}={:.3f}'.format(
  292.             epoch, (time.time()-tic), name1, loss1, name2, loss2, name3, loss3, name4, loss4))
  293.         if not (epoch + 1) % args.val_interval:
  294.             print("validate:", epoch + 1)
  295.             # consider reduce the frequency of validation to save time
  296.             #map_name, mean_ap = validate(net, val_data, ctx, eval_metric)
  297.             #val_msg = '\n'.join(['{}={}'.format(k, v) for k, v in zip(map_name, mean_ap)])
  298.             #logger.info('[Epoch {}] Validation: \n{}'.format(epoch, val_msg))
  299.             #current_map = float(mean_ap[-1])
  300.             current_map = 0.
  301.         else:
  302.             current_map = 0.
  303.         #save_params(net, best_map, current_map, epoch, args.save_interval, args.save_prefix)
  304.         net.save_parameters('yolo3_pikachu')
  305.     net.save_parameters('yolo3_pikachu')
  306.  
  307. if __name__ == '__main__':
  308.     args = parse_args()
  309.     # fix seed for mxnet, numpy and python builtin random generator.
  310.     gutils.random.seed(args.seed)
  311.  
  312.     # training contexts
  313.     ctx = [mx.gpu(int(i)) for i in args.gpus.split(',') if i.strip()]
  314.     ctx = ctx if ctx else [mx.cpu()]
  315.  
  316.     # network
  317.     net_name = args.network
  318.     args.save_prefix += net_name
  319.     # use sync bn if specified
  320.     num_sync_bn_devices = len(ctx) if args.syncbn else -1
  321.     classes = read_classes(args)
  322.     net = None
  323.     if num_sync_bn_devices > 1:
  324.         print("num_sync_bn_devices > 1")
  325.         if args.pretrained == 0:
  326.             net = get_model(net_name, pretrained=True, num_sync_bn_devices=num_sync_bn_devices)
  327.         else:        
  328.             net = get_model(net_name, pretrained_base=True, num_sync_bn_devices=num_sync_bn_devices)
  329.            
  330.         net.reset_class(classes)            
  331.         async_net = get_model(net_name, pretrained_base=False)  # used by cpu worker
  332.     else:
  333.         print("num_sync_bn_devices <= 1")        
  334.         if args.pretrained == 0:
  335.             net = get_model(net_name, pretrained=True)            
  336.         else:
  337.             net = get_model(net_name, pretrained_base=True)
  338.         net.reset_class(classes)            
  339.         async_net = net
  340.            
  341.     if args.resume.strip():
  342.         net.load_parameters(args.resume.strip())
  343.         async_net.load_parameters(args.resume.strip())
  344.     else:
  345.         with warnings.catch_warnings(record=True) as w:
  346.             warnings.simplefilter("always")
  347.             net.initialize()
  348.             async_net.initialize()
  349.  
  350.     # training data
  351.     train_dataset, val_dataset, eval_metric = get_dataset(args)
  352.     train_data, val_data = get_dataloader(
  353.         async_net, train_dataset, val_dataset, args.data_shape, args.batch_size, args.num_workers, args)
  354.  
  355.     # training
  356.     #train(net, train_data, val_data, eval_metric, ctx, args)
  357.    
  358.     test_url = 'https://raw.githubusercontent.com/zackchase/mxnet-the-straight-dope/master/img/pikachu.jpg'
  359.     download(test_url, 'pikachu_test.jpg')
  360.     net = gcv.model_zoo.get_model('yolo3_darknet53_custom', classes=classes, pretrained_base=False)
  361.     net.load_parameters('yolo3_pikachu')
  362.     net.collect_params().reset_ctx(ctx)
  363.     validate(net, val_data, ctx, eval_metric)
  364.     #x, image = gcv.data.transforms.presets.yolo.load_test('pikachu_test.jpg')
  365.     #cid, score, bbox = net(x)
  366.     #ax = viz.plot_bbox(image, bbox[0], score[0], cid[0], class_names=classes)
  367.     #plt.show()
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top