diff --git a/embedding_converter/README.md b/embedding_converter/README.md index e74162e..873b2e3 100644 --- a/embedding_converter/README.md +++ b/embedding_converter/README.md @@ -69,6 +69,7 @@ file_pattern = arcface_converter_simswap_{epoch:02d}_{val_loss:.4f} directory_path = .exports source_path = .outputs/last.ckpt target_path = .exports/arcface_converter_simswap.onnx +ir_version = 10 opset_version = 15 ``` diff --git a/embedding_converter/config.ini b/embedding_converter/config.ini index 4bda383..71166fd 100644 --- a/embedding_converter/config.ini +++ b/embedding_converter/config.ini @@ -29,6 +29,7 @@ file_pattern = directory_path = source_path = target_path = +ir_version = opset_version = [execution] diff --git a/embedding_converter/src/exporting.py b/embedding_converter/src/exporting.py index bb01e17..8cb0532 100644 --- a/embedding_converter/src/exporting.py +++ b/embedding_converter/src/exporting.py @@ -13,10 +13,12 @@ def export() -> None: directory_path = CONFIG.get('exporting', 'directory_path') source_path = CONFIG.get('exporting', 'source_path') target_path = CONFIG.get('exporting', 'target_path') + ir_version = CONFIG.getint('exporting', 'ir_version') opset_version = CONFIG.getint('exporting', 'opset_version') makedirs(directory_path, exist_ok = True) - embedding_converter_trainer = EmbeddingConverterTrainer.load_from_checkpoint(source_path, map_location = 'cpu') - embedding_converter_trainer.eval() + model = EmbeddingConverterTrainer.load_from_checkpoint(source_path, map_location = 'cpu') + model.eval() + model.ir_version = ir_version input_tensor = torch.randn(1, 512) - torch.onnx.export(embedding_converter_trainer, input_tensor, target_path, input_names = [ 'input' ], output_names = [ 'output' ], opset_version = opset_version) + torch.onnx.export(model, input_tensor, target_path, input_names = [ 'input' ], output_names = [ 'output' ], opset_version = opset_version) diff --git a/face_swapper/README.md b/face_swapper/README.md index 332b175..da0618d 100644 --- a/face_swapper/README.md +++ b/face_swapper/README.md @@ -92,6 +92,7 @@ validation_frequency = 1000 directory_path = .exports source_path = .outputs/last.ckpt target_path = .exports/face_swapper.onnx +ir_version = 10 opset_version = 15 ``` diff --git a/face_swapper/config.ini b/face_swapper/config.ini index faa6254..01e343f 100644 --- a/face_swapper/config.ini +++ b/face_swapper/config.ini @@ -48,6 +48,7 @@ validation_frequency = directory_path = source_path = target_path = +ir_version = opset_version = [inferencing] diff --git a/face_swapper/src/exporting.py b/face_swapper/src/exporting.py index 311fdc8..d4c12ec 100644 --- a/face_swapper/src/exporting.py +++ b/face_swapper/src/exporting.py @@ -13,6 +13,7 @@ def export() -> None: directory_path = CONFIG.get('exporting', 'directory_path') source_path = CONFIG.get('exporting', 'source_path') target_path = CONFIG.get('exporting', 'target_path') + ir_version = CONFIG.getint('exporting', 'ir_version') opset_version = CONFIG.getint('exporting', 'opset_version') makedirs(directory_path, exist_ok = True) @@ -20,6 +21,7 @@ def export() -> None: model = AdaptiveEmbeddingIntegrationNetwork() model.load_state_dict(state_dict) model.eval() + model.ir_version = ir_version source_tensor = torch.randn(1, 512) target_tensor = torch.randn(1, 3, 256, 256) - torch.onnx.export(model, (target_tensor, source_tensor), target_path, input_names = [ 'target', 'source' ], output_names = [ 'output' ], opset_version = opset_version) + torch.onnx.export(model, (source_tensor, target_tensor), target_path, input_names = [ 'source', 'target' ], output_names = [ 'output' ], opset_version = opset_version)