def save_root(args, output_path, data_config, scores, labels, observers): """ Saves as .root :param data_config: :param scores: :param labels :param observers :return: """ from utils.data.fileio import _write_root output = {} if args.regression_mode: output[data_config.label_names[0]] = labels[data_config.label_names[0]] output['output'] = scores else: for idx, label_name in enumerate(data_config.label_value): output[label_name] = (labels[data_config.label_names[0]] == idx) output['score_' + label_name] = scores[:, idx] for k, v in labels.items(): if k == data_config.label_names[0]: continue if v.ndim > 1: _logger.warning('Ignoring %s, not a 1d array.', k) continue output[k] = v for k, v in observers.items(): if v.ndim > 1: _logger.warning('Ignoring %s, not a 1d array.', k) continue output[k] = v _write_root(output_path, output)
def main(): parser = argparse.ArgumentParser() parser.add_argument('-c', '--data-config', type=str, default='data/ak15_points_pf_sv_v0.yaml', help='data config YAML file') parser.add_argument('-i', '--data-train', nargs='*', default=[], help='training files') parser.add_argument('-t', '--data-test', nargs='*', default=[], help='testing files') parser.add_argument( '--data-fraction', type=float, default=1, help= 'fraction of events to load from each file; for training, the events are randomly selected for each epoch' ) parser.add_argument( '--data-dilation', type=int, default=1, help= 'reduce number of file by a factor of `d` for training. NOT recommended in general - use `--data-fraction` instead.' ) parser.add_argument( '--files-per-fetch', type=int, default=20, help= 'number of files to load each time; shuffling is done within these events, so choose a number large enough to get events from all classes' ) parser.add_argument('--train-val-split', type=float, default=0.8, help='training/validation split fraction') parser.add_argument( '--demo', action='store_true', default=False, help= 'quickly test the setup by running over only a small number of events') parser.add_argument( '--lr-finder', type=str, default=None, help= 'run learning rate finder instead of the actual training; format: ``start_lr, end_lr, num_iters``' ) parser.add_argument( '-n', '--network-config', type=str, default='networks/particle_net_pfcand_sv.py', help= 'network architecture configuration file; the path must be relative to the current dir' ) parser.add_argument( '--network-option', nargs=2, action='append', default=[], help= 'options to pass to the model class constructor, e.g., `--network-option use_counts False`' ) parser.add_argument( '-m', '--model-prefix', type=str, default='test_output/model_name', help= 'path to save or load the model; for training, this will be used as a prefix; for testing, this should be the full path including extension' ) parser.add_argument('--num-epochs', type=int, default=20, help='number of epochs') parser.add_argument( '--optimizer', type=str, default='ranger', choices=['adam', 'ranger'], # TODO: add more help='optimizer for the training') parser.add_argument( '--load-epoch', type=int, default=None, help= 'used to resume interrupted training, load model and optimizer state saved in the `epoch-%d_state.pt` and `epoch-%d_optimizer.pt` files' ) parser.add_argument('--start-lr', type=float, default=5e-3, help='start learning rate') parser.add_argument( '--lr-steps', type=str, default='10,20', help= 'steps to reduce the lr; currently only used when setting `--optimizer` to adam' ) parser.add_argument('--batch-size', type=int, default=128, help='batch size') parser.add_argument( '--use-amp', action='store_true', default=False, help='use mixed precision training (fp16); NOT WORKING YET') parser.add_argument( '--gpus', type=str, default='0', help='device for the training/testing; to use CPU, set to empty string (' '); to use multiple gpu, set it as a comma separated list, e.g., `1,2,3,4`' ) parser.add_argument( '--num-workers', type=int, default=2, help= 'number of threads to load the dataset; memory consuption and disk access load increases (~linearly) with this numbers' ) parser.add_argument('--predict', action='store_true', default=False, help='run prediction instead of training') parser.add_argument( '--predict-output', type=str, help= 'path to save the prediction output, support `.root` and `.awkd` format' ) parser.add_argument( '--export-onnx', type=str, default=None, help= 'export the PyTorch model to ONNX model and save it at the given path (path must ends w/ .onnx); ' 'needs to set `--data-config`, `--network-config`, and `--model-prefix` (requires the full model path)' ) args = parser.parse_args() _logger.info(args) if args.use_amp: raise NotImplementedError # from apex import amp if args.data_dilation > 1: _logger.warning( 'Use of `data-dilation` is not recomended in general -- consider using `data-fraction` instead.' ) # training/testing mode training_mode = not args.predict # device if args.gpus: gpus = [int(i) for i in args.gpus.split(',')] dev = torch.device(gpus[0]) else: gpus = None dev = torch.device('cpu') # load data if training_mode: filelist = sorted(sum([glob.glob(f) for f in args.data_train], [])) # np.random.seed(1) np.random.shuffle(filelist) if args.demo: filelist = filelist[:20] _logger.info(filelist) args.data_fraction = 0.1 args.files_per_fetch = 5 train_data = SimpleIterDataset(filelist, args.data_config, for_training=True, partial_load=((0, args.train_val_split), args.data_fraction), dilation=args.data_dilation, files_per_fetch=args.files_per_fetch) val_data = SimpleIterDataset(filelist, args.data_config, for_training=True, partial_load=((args.train_val_split, 1), args.data_fraction), dilation=args.data_dilation, files_per_fetch=args.files_per_fetch) train_loader = DataLoader(train_data, num_workers=args.num_workers, batch_size=args.batch_size, drop_last=True, pin_memory=True) val_loader = DataLoader(val_data, num_workers=args.num_workers, batch_size=args.batch_size, drop_last=True, pin_memory=True) data_config = train_data.config else: filelist = sorted(sum([glob.glob(f) for f in args.data_test], [])) test_data = SimpleIterDataset(filelist, args.data_config, for_training=False, files_per_fetch=1) test_loader = DataLoader(test_data, num_workers=args.num_workers, batch_size=args.batch_size, drop_last=False, pin_memory=True) data_config = test_data.config # model network_module = import_module( args.network_config.replace('.py', '').replace('/', '.')) network_options = {k: ast.literal_eval(v) for k, v in args.network_option} if args.export_onnx: network_options['for_inference'] = True model, model_info = network_module.get_model(data_config, **network_options) _logger.info(model) # export to ONNX if args.export_onnx: assert (args.export_onnx.endswith('.onnx')) model_path = args.model_prefix _logger.info('Exporting model %s to ONNX' % model_path) model.load_state_dict(torch.load(model_path, map_location='cpu')) model = model.cpu() model.eval() os.makedirs(os.path.dirname(args.export_onnx), exist_ok=True) inputs = tuple( torch.ones(model_info['input_shapes'][k], dtype=torch.float32) for k in model_info['input_names']) torch.onnx.export(model, inputs, args.export_onnx, input_names=model_info['input_names'], output_names=model_info['output_names'], dynamic_axes=model_info.get('dynamic_axes', None), opset_version=11) _logger.info('ONNX model saved to %s', args.export_onnx) return # note: we should always save/load the state_dict of the original model, not the one wrapped by nn.DataParallel # so we do not convert it to nn.DataParallel now model = model.to(dev) # loss function try: loss_func = network_module.get_loss(data_config, **network_options) _logger.info(loss_func) except AttributeError: loss_func = torch.nn.CrossEntropyLoss() _logger.warning( 'Loss function not defined in %s. Will use `torch.nn.CrossEntropyLoss()` by default.', args.network_config) if training_mode: # optimizer & learning rate if args.optimizer == 'adam': opt = torch.optim.Adam(model.parameters(), lr=args.start_lr) if args.lr_finder is None: lr_steps = [int(x) for x in args.lr_steps.split(',')] scheduler = torch.optim.lr_scheduler.MultiStepLR( opt, milestones=lr_steps, gamma=0.1) else: from utils.nn.optimizer.ranger import Ranger opt = Ranger(model.parameters(), lr=args.start_lr) if args.lr_finder is None: lr_decay_epochs = max(1, int(args.num_epochs * 0.3)) lr_decay_rate = 0.01**(1. / lr_decay_epochs) scheduler = torch.optim.lr_scheduler.MultiStepLR( opt, milestones=list( range(args.num_epochs - lr_decay_epochs, args.num_epochs)), gamma=lr_decay_rate) # TODO: mixed precision training if args.use_amp: # model, opt = amp.initialize( # model, opt, opt_level="O2", # keep_batchnorm_fp32=True, loss_scale="dynamic" # ) model, opt = amp.initialize(model, opt, opt_level="O1", keep_batchnorm_fp32=None, loss_scale="dynamic") # load previous training and resume if `--load-epoch` is set if args.load_epoch is not None: _logger.info('Resume training from epoch %d' % args.load_epoch) model_state = torch.load(args.model_prefix + '_epoch-%d_state.pt' % args.load_epoch, map_location=dev) model.load_state_dict(model_state) opt_state = torch.load(args.model_prefix + '_epoch-%d_optimizer.pt' % args.load_epoch, map_location=dev) opt.load_state_dict(opt_state) # mutli-gpu if gpus is not None and len(gpus) > 1: model = torch.nn.DataParallel( model, device_ids=gpus ) # model becomes `torch.nn.DataParallel` w/ model.module being the orignal `torch.nn.Module` model = model.to(dev) # lr finder: keep it after all other setups if args.lr_finder is not None: start_lr, end_lr, num_iter = args.lr_finder.replace(' ', '').split(',') from utils.lr_finder import LRFinder lr_finder = LRFinder(model, opt, loss_func, device=dev, input_names=train_data.config.input_names, label_names=train_data.config.label_names) lr_finder.range_test(train_loader, start_lr=float(start_lr), end_lr=float(end_lr), num_iter=int(num_iter)) lr_finder.plot(output='lr_finder.png' ) # to inspect the loss-learning rate graph return # training loop best_valid_acc = 0 for epoch in range(args.num_epochs): if args.load_epoch is not None: if epoch <= args.load_epoch: continue print('-' * 50) _logger.info('Epoch #%d training' % epoch) train(model, loss_func, opt, scheduler, train_loader, dev) if args.model_prefix: dirname = os.path.dirname(args.model_prefix) if dirname and not os.path.exists(dirname): os.makedirs(dirname) state_dict = model.module.state_dict() if isinstance( model, torch.nn.DataParallel) else model.state_dict() torch.save(state_dict, args.model_prefix + '_epoch-%d_state.pt' % epoch) torch.save( opt.state_dict(), args.model_prefix + '_epoch-%d_optimizer.pt' % epoch) _logger.info('Epoch #%d validating' % epoch) valid_acc = evaluate(model, val_loader, dev, loss_func=loss_func) if valid_acc > best_valid_acc: best_valid_acc = valid_acc if args.model_prefix: shutil.copy2( args.model_prefix + '_epoch-%d_state.pt' % epoch, args.model_prefix + '_best_acc_state.pt') torch.save(model, args.model_prefix + '_best_acc_full.pt') _logger.info( 'Epoch #%d: Current validation acc: %.5f (best: %.5f)' % (epoch, valid_acc, best_valid_acc)) else: # run prediction if args.model_prefix.endswith('.onnx'): _logger.info('Loading model %s for eval' % args.model_prefix) from utils.nn.tools import evaluate_onnx test_acc, scores, labels, observers = evaluate_onnx( args.model_prefix, test_loader) else: model_path = args.model_prefix if args.model_prefix.endswith( '.pt') else args.model_prefix + '_best_acc_state.pt' _logger.info('Loading model %s for eval' % model_path) model.load_state_dict(torch.load(model_path, map_location=dev)) if gpus is not None and len(gpus) > 1: model = torch.nn.DataParallel(model, device_ids=gpus) model = model.to(dev) test_acc, scores, labels, observers = evaluate(model, test_loader, dev, for_training=False) _logger.info('Test acc %.5f' % test_acc) if args.predict_output: os.makedirs(os.path.dirname(args.predict_output), exist_ok=True) if args.predict_output.endswith('.root'): from utils.data.fileio import _write_root output = {} for idx, label_name in enumerate(data_config.label_value): output[label_name] = ( labels[data_config.label_names[0]] == idx) output['score_' + label_name] = scores[:, idx] for k, v in labels.items(): if k == data_config.label_names[0]: continue if v.ndim > 1: _logger.warning('Ignoring %s, not a 1d array.', k) continue output[k] = v for k, v in observers.items(): if v.ndim > 1: _logger.warning('Ignoring %s, not a 1d array.', k) continue output[k] = v _write_root(args.predict_output, output) else: import awkward output = {'scores': scores} output.update(labels) output.update(observers) awkward.save(args.predict_output, output, mode='w') _logger.info('Written output to %s' % args.predict_output)
def main(args): _logger.info('args:\n - %s', '\n - '.join(str(it) for it in args.__dict__.items())) if args.file_fraction < 1: _logger.warning( 'Use of `file-fraction` is not recommended in general -- prefer using `data-fraction` instead.' ) # classification/regression mode if args.regression_mode: _logger.info('Running in regression mode') from utils.nn.tools import train_regression as train from utils.nn.tools import evaluate_regression as evaluate else: _logger.info('Running in classification mode') from utils.nn.tools import train_classification as train from utils.nn.tools import evaluate_classification as evaluate # training/testing mode training_mode = not args.predict # device if args.gpus: gpus = [int(i) for i in args.gpus.split(',')] dev = torch.device(gpus[0]) else: gpus = None dev = torch.device('cpu') # load data if training_mode: train_loader, val_loader, data_config, train_input_names, train_label_names = train_load( args) else: test_loaders, data_config = test_load(args) if args.io_test: data_loader = train_loader if training_mode else list( test_loaders.values())[0]() iotest(args, data_loader) return model, model_info, network_module, network_options = model_setup( args, data_config) if args.print: return if args.profile: profile(args, model, model_info, device=dev) return # export to ONNX if args.export_onnx: onnx(args, model, data_config, model_info) return if args.tensorboard: from utils.nn.tools import TensorboardHelper tb = TensorboardHelper(tb_comment=args.tensorboard, tb_custom_fn=args.tensorboard_custom_fn) else: tb = None # note: we should always save/load the state_dict of the original model, not the one wrapped by nn.DataParallel # so we do not convert it to nn.DataParallel now orig_model = model if training_mode: model = orig_model.to(dev) # loss function try: loss_func = network_module.get_loss(data_config, **network_options) _logger.info('Using loss function %s with options %s' % (loss_func, network_options)) except AttributeError: loss_func = torch.nn.CrossEntropyLoss() _logger.warning( 'Loss function not defined in %s. Will use `torch.nn.CrossEntropyLoss()` by default.', args.network_config) # optimizer & learning rate opt, scheduler = optim(args, model, dev) # multi-gpu if gpus is not None and len(gpus) > 1: # model becomes `torch.nn.DataParallel` w/ model.module being the original `torch.nn.Module` model = torch.nn.DataParallel(model, device_ids=gpus) model = model.to(dev) # lr finder: keep it after all other setups if args.lr_finder is not None: start_lr, end_lr, num_iter = args.lr_finder.replace(' ', '').split(',') from utils.lr_finder import LRFinder lr_finder = LRFinder(model, opt, loss_func, device=dev, input_names=train_input_names, label_names=train_label_names) lr_finder.range_test(train_loader, start_lr=float(start_lr), end_lr=float(end_lr), num_iter=int(num_iter)) lr_finder.plot(output='lr_finder.png' ) # to inspect the loss-learning rate graph return if args.use_amp: from torch.cuda.amp import GradScaler scaler = GradScaler() else: scaler = None # training loop best_valid_metric = np.inf if args.regression_mode else 0 for epoch in range(args.num_epochs): if args.load_epoch is not None: if epoch <= args.load_epoch: continue print('-' * 50) _logger.info('Epoch #%d training' % epoch) train(model, loss_func, opt, scheduler, train_loader, dev, epoch, steps_per_epoch=args.steps_per_epoch, grad_scaler=scaler, tb_helper=tb) if args.model_prefix: dirname = os.path.dirname(args.model_prefix) if dirname and not os.path.exists(dirname): os.makedirs(dirname) state_dict = model.module.state_dict() if isinstance( model, torch.nn.DataParallel) else model.state_dict() torch.save(state_dict, args.model_prefix + '_epoch-%d_state.pt' % epoch) torch.save( opt.state_dict(), args.model_prefix + '_epoch-%d_optimizer.pt' % epoch) _logger.info('Epoch #%d validating' % epoch) valid_metric = evaluate(model, val_loader, dev, epoch, loss_func=loss_func, steps_per_epoch=args.steps_per_epoch_val, tb_helper=tb) is_best_epoch = (valid_metric < best_valid_metric ) if args.regression_mode else ( valid_metric > best_valid_metric) if is_best_epoch: best_valid_metric = valid_metric if args.model_prefix: shutil.copy2( args.model_prefix + '_epoch-%d_state.pt' % epoch, args.model_prefix + '_best_epoch_state.pt') torch.save(model, args.model_prefix + '_best_epoch_full.pt') _logger.info( 'Epoch #%d: Current validation metric: %.5f (best: %.5f)' % (epoch, valid_metric, best_valid_metric), color='bold') if args.data_test: if training_mode: del train_loader, val_loader test_loaders, data_config = test_load(args) if not args.model_prefix.endswith('.onnx'): model = orig_model.to(dev) model_path = args.model_prefix if args.model_prefix.endswith( '.pt') else args.model_prefix + '_best_epoch_state.pt' _logger.info('Loading model %s for eval' % model_path) model.load_state_dict(torch.load(model_path, map_location=dev)) if gpus is not None and len(gpus) > 1: model = torch.nn.DataParallel(model, device_ids=gpus) model = model.to(dev) for name, get_test_loader in test_loaders.items(): test_loader = get_test_loader() # run prediction if args.model_prefix.endswith('.onnx'): _logger.info('Loading model %s for eval' % args.model_prefix) from utils.nn.tools import evaluate_onnx test_metric, scores, labels, observers = evaluate_onnx( args.model_prefix, test_loader) else: test_metric, scores, labels, observers = evaluate( model, test_loader, dev, epoch=None, for_training=False, tb_helper=tb) _logger.info('Test metric %.5f' % test_metric, color='bold') del test_loader if args.predict_output: if '/' not in args.predict_output: args.predict_output = os.path.join( os.path.dirname(args.model_prefix), 'predict_output', args.predict_output) os.makedirs(os.path.dirname(args.predict_output), exist_ok=True) if name == '': output_path = args.predict_output else: base, ext = os.path.splitext(args.predict_output) output_path = base + '_' + name + ext if output_path.endswith('.root'): save_root(args, output_path, data_config, scores, labels, observers) else: save_awk(args, output_path, scores, labels, observers) _logger.info('Written output to %s' % output_path, color='bold')
def main(args): _logger.info(args) if args.file_fraction < 1: _logger.warning( 'Use of `file-fraction` is not recommended in general -- prefer using `data-fraction` instead.' ) # training/testing mode training_mode = not args.predict # device if args.gpus: gpus = [int(i) for i in args.gpus.split(',')] dev = torch.device(gpus[0]) else: gpus = None dev = torch.device('cpu') # load data if training_mode: train_loader, val_loader, data_config, train_input_names, train_label_names = train_load( args) else: test_loader, data_config = test_load(args) if args.io_test: data_loader = train_loader if training_mode else test_loader iotest(args, data_loader) return model, model_info, network_module = _model(args, data_config) # export to ONNX if args.export_onnx: onnx(args, model, data_config, model_info) return # note: we should always save/load the state_dict of the original model, not the one wrapped by nn.DataParallel # so we do not convert it to nn.DataParallel now model = model.to(dev) if training_mode: # loss function try: loss_func = network_module.get_loss(data_config, **network_options) _logger.info(loss_func) except AttributeError: loss_func = torch.nn.CrossEntropyLoss() _logger.warning( 'Loss function not defined in %s. Will use `torch.nn.CrossEntropyLoss()` by default.', args.network_config) # optimizer & learning rate opt, scheduler = optim(args, model) # load previous training and resume if `--load-epoch` is set if args.load_epoch is not None: _logger.info('Resume training from epoch %d' % args.load_epoch) model_state = torch.load(args.model_prefix + '_epoch-%d_state.pt' % args.load_epoch, map_location=dev) model.load_state_dict(model_state) opt_state = torch.load(args.model_prefix + '_epoch-%d_optimizer.pt' % args.load_epoch, map_location=dev) opt.load_state_dict(opt_state) # multi-gpu if gpus is not None and len(gpus) > 1: model = torch.nn.DataParallel( model, device_ids=gpus ) # model becomes `torch.nn.DataParallel` w/ model.module being the orignal `torch.nn.Module` model = model.to(dev) # lr finder: keep it after all other setups if args.lr_finder is not None: start_lr, end_lr, num_iter = args.lr_finder.replace(' ', '').split(',') from utils.lr_finder import LRFinder lr_finder = LRFinder(model, opt, loss_func, device=dev, input_names=train_input_names, label_names=train_label_names) lr_finder.range_test(train_loader, start_lr=float(start_lr), end_lr=float(end_lr), num_iter=int(num_iter)) lr_finder.plot(output='lr_finder.png' ) # to inspect the loss-learning rate graph return if args.use_amp: from torch.cuda.amp import GradScaler scaler = GradScaler() else: scaler = None # training loop best_valid_acc = 0 for epoch in range(args.num_epochs): if args.load_epoch is not None: if epoch <= args.load_epoch: continue print('-' * 50) _logger.info('Epoch #%d training' % epoch) train(model, loss_func, opt, scheduler, train_loader, dev, grad_scaler=scaler) if args.model_prefix: dirname = os.path.dirname(args.model_prefix) if dirname and not os.path.exists(dirname): os.makedirs(dirname) state_dict = model.module.state_dict() if isinstance( model, torch.nn.DataParallel) else model.state_dict() torch.save(state_dict, args.model_prefix + '_epoch-%d_state.pt' % epoch) torch.save( opt.state_dict(), args.model_prefix + '_epoch-%d_optimizer.pt' % epoch) _logger.info('Epoch #%d validating' % epoch) valid_acc = evaluate(model, val_loader, dev, loss_func=loss_func) if valid_acc > best_valid_acc: best_valid_acc = valid_acc if args.model_prefix: shutil.copy2( args.model_prefix + '_epoch-%d_state.pt' % epoch, args.model_prefix + '_best_acc_state.pt') torch.save(model, args.model_prefix + '_best_acc_full.pt') _logger.info( 'Epoch #%d: Current validation acc: %.5f (best: %.5f)' % (epoch, valid_acc, best_valid_acc)) else: # run prediction predict_model(args, test_loader, model, dev, data_config, gpus)