()
| 15 | import tza |
| 16 | |
| 17 | def main(): |
| 18 | # Parse the command line arguments |
| 19 | cfg = parse_args(description='Exports a trained model to the runtime weights (TZA) or some other format.') |
| 20 | |
| 21 | print('Result:', cfg.result) |
| 22 | |
| 23 | if cfg.target == 'package': |
| 24 | # Get the output filename |
| 25 | if cfg.output: |
| 26 | output_filename = cfg.output |
| 27 | else: |
| 28 | output_filename = os.path.join(cfg.results_dir, cfg.result) + '.zip' |
| 29 | print('Output:', output_filename) |
| 30 | |
| 31 | # Get the list of files that belong to the result (latest checkpoint only) |
| 32 | result_dir = get_result_dir(cfg) |
| 33 | filenames = [get_config_filename(result_dir)] |
| 34 | filenames.append(get_checkpoint_state_filename(result_dir)) |
| 35 | latest_epoch = get_latest_checkpoint_epoch(result_dir) |
| 36 | filenames.append(get_checkpoint_filename(result_dir, latest_epoch)) |
| 37 | filenames += glob(os.path.join(get_result_log_dir(result_dir), 'events.out.*')) |
| 38 | filenames += glob(os.path.join(result_dir, 'src.*')) |
| 39 | |
| 40 | # Save the ZIP file |
| 41 | save_zip(output_filename, filenames, root_dir=cfg.results_dir) |
| 42 | else: |
| 43 | # Initialize the PyTorch device |
| 44 | device = init_device(cfg) |
| 45 | |
| 46 | # Load the result config |
| 47 | result_dir = get_result_dir(cfg) |
| 48 | if not os.path.isdir(result_dir): |
| 49 | error('result does not exist') |
| 50 | result_cfg = load_config(result_dir) |
| 51 | |
| 52 | # Initialize the model |
| 53 | if cfg.target in {'onnx', 'onnx_noparams'}: |
| 54 | model = get_model(result_cfg) |
| 55 | model.to(device) |
| 56 | else: |
| 57 | model = None |
| 58 | |
| 59 | # Load the checkpoint |
| 60 | checkpoint = load_checkpoint(result_dir, device, cfg.num_epochs, model) |
| 61 | epoch = checkpoint['epoch'] |
| 62 | model_state = checkpoint['model_state'] |
| 63 | print('Epoch:', epoch) |
| 64 | |
| 65 | if cfg.target == 'weights': |
| 66 | # Save the weights to a TZA file |
| 67 | if cfg.output: |
| 68 | output_filename = cfg.output |
| 69 | else: |
| 70 | output_filename = os.path.join(result_dir, cfg.result) |
| 71 | if cfg.num_epochs: |
| 72 | output_filename += '_%d' % epoch |
| 73 | output_filename += '.tza' |
| 74 | print('Output:', output_filename) |
no test coverage detected