def save_model_training( model: Module, optim: Optimizer, input_shape: Tuple[int, ...], save_name: str, save_dir: str, epoch: int, val_res: Union[ModuleRunResults, None], convert_qat: bool = False, ): """ :param model: model architecture :param optim: The optimizer used :param input_shape: A tuple of integers representing the input shape :param save_name: name to save model to :param save_dir: directory to save results in :param epoch: integer representing umber of epochs to :param val_res: results from validation run :param convert_qat: True if model is to be quantized before saving """ has_top1 = "top1acc" in val_res.results metric_name = "top-1 accuracy" if has_top1 else "val_loss" metric = val_res.result_mean( "top1acc" if has_top1 else DEFAULT_LOSS_KEY).item() print(f"Saving model for epoch {epoch} and {metric_name} " f"{metric} to {save_dir} for {save_name}") exporter = ModuleExporter(model, save_dir) exporter.export_pytorch(optim, epoch, f"{save_name}.pth") exporter.export_onnx( torch.randn(1, *input_shape), f"{save_name}.onnx", convert_qat=convert_qat, ) info_path = os.path.join(save_dir, f"{save_name}.txt") with open(info_path, "w") as info_file: info_lines = [ f"epoch: {epoch}", ] if val_res is not None: for loss in val_res.results.keys(): info_lines.append( f"{loss}: {val_res.result_mean(loss).item()}") info_file.write("\n".join(info_lines))
def export(args: ExportArgs, model: Module, val_loader: DataLoader, save_dir: str) -> None: """ Utility method to export the model and data :param args : An ExportArgs object containing config for export task. :param model: loaded model architecture to export :param val_loader: A DataLoader for validation data :param save_dir: Directory to store checkpoints at during exporting process """ exporter = ModuleExporter(model, save_dir) # export PyTorch state dict LOGGER.info(f"exporting pytorch in {save_dir}") exporter.export_pytorch(use_zipfile_serialization_if_available=( args.use_zipfile_serialization_if_available)) onnx_exported = False for batch, data in tqdm( enumerate(val_loader), desc="Exporting samples", total=args.num_samples if args.num_samples > 1 else 1, ): if not onnx_exported: # export onnx file using first sample for graph freezing LOGGER.info(f"exporting onnx in {save_dir}") exporter.export_onnx(data[0], opset=args.onnx_opset, convert_qat=True) onnx_exported = True if args.num_samples > 0: exporter.export_samples(sample_batches=[data[0]], sample_labels=[data[1]], exp_counter=batch)
def main(args): ############################ # logging and saving setup # ############################ save_dir = os.path.abspath(os.path.expanduser(args.save_dir)) # get unique model tag, defaults to '{model_name}' if not args.model_tag: model_tag = args.model.replace("/", ".") model_id = model_tag model_inc = 0 while os.path.exists(os.path.join(args.save_dir, model_id)): model_inc += 1 model_id = "{}__{:02d}".format(model_tag, model_inc) else: model_id = args.model_tag save_dir = os.path.join(save_dir, model_id) create_dirs(save_dir) print("Model id is set to {}".format(model_id)) ########################### # standard training setup # ########################### # create data loaders train_loader, _, _ = _create_imagefolder_dataloader(args, train=True) val_loader, num_classes, image_shape = _create_imagefolder_dataloader( args, train=False ) dataloaders = {"train": train_loader, "val": val_loader} # create model model = _get_torchvision_model( args.model, num_classes, args.pretrained, args.checkpoint_path, ) print("created model: {}".format(model)) device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) print("using device: {}".format(device)) # create standard SGD optimizer and cross entropy loss function criterion = CrossEntropyLoss() optimizer = SGD( model.parameters(), lr=0.001, momentum=0.9 ) # lr will be overridden by recipe ########################## # add sparseml modifiers # ########################## manager = ScheduledModifierManager.from_yaml(args.recipe_path) optimizer = ScheduledOptimizer( optimizer, model, manager, steps_per_epoch=len(train_loader), loggers=[PythonLogger()], ) ######################## # torchvision training # ######################## model, val_acc_history = train_model( model, dataloaders, criterion, optimizer, device, num_epochs=manager.max_epochs, is_inception="inception" in args.model, ) ######################## # export trained model # ######################## exporter = ModuleExporter(model, save_dir) sample_input = torch.randn(image_shape).unsqueeze(0) # sample batch for ONNX export exporter.export_onnx(sample_input) exporter.export_pytorch() print("Model ONNX export and PyTorch weights saved to {}".format(save_dir))
def train( working_dir: str, config_path: str, model: Module, train_dataset: Dataset, val_dataset: Dataset, batch_size: int, optim_const: Callable[[Module], Optimizer], loss: Union[LossWrapper, Callable[[Any, Any], Tensor]], devices: str, ): """ Dataset setup """ LOGGER.info("batch_size set to {}".format(batch_size)) LOGGER.info("train_dataset set to {}".format(train_dataset)) LOGGER.info("val_dataset set to {}".format(val_dataset)) train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True, ) val_loader = DataLoader( val_dataset, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True, ) """ Model, optimizer, loss setup """ model_dir = clean_path(os.path.join(working_dir, "model")) optim = optim_const(model) LOGGER.info("model set to {}".format(model)) LOGGER.info("optimizer set to {}".format(optim)) LOGGER.info("loss set to {}".format(loss)) LOGGER.info("devices set to {}".format(devices)) """ Manager and config setup """ manager = ScheduledModifierManager.from_yaml(config_path) logs_dir = clean_path(os.path.join(working_dir, "logs")) loggers = [TensorBoardLogger(logs_dir), PythonLogger()] optim = ScheduledOptimizer( optim, model, manager, steps_per_epoch=len(train_loader), loggers=loggers ) """ Training and testing """ model, device, device_ids = model_to_device(model, devices) trainer = ModuleTrainer(model, device, loss, optim, loggers=loggers) tester = ModuleTester(model, device, loss, loggers=loggers, log_steps=-1) epoch = -1 tester.run_epoch(val_loader, epoch=epoch) for epoch in range(manager.max_epochs): LOGGER.info("starting training epoch {}".format(epoch)) train_res = trainer.run_epoch(train_loader, epoch) LOGGER.info("finished training epoch {}: {}".format(epoch, train_res)) val_res = tester.run_epoch(val_loader, epoch) LOGGER.info("finished validation epoch {}: {}".format(epoch, val_res)) exporter = ModuleExporter(model, model_dir) exporter.export_pytorch(optim, epoch) for data in val_loader: exporter.export_onnx(data)