Skip to content

Commit 06d1bfc

Browse files
mgumowskmgumowsk
authored
[poc] pytorch model converter tool (#421)
* tools/model_converter * fix bandit * checks * fix * style * semgrep * sem
1 parent 6b54085 commit 06d1bfc

File tree

3 files changed

+837
-0
lines changed

3 files changed

+837
-0
lines changed

‎tools/model_converter/README.md‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
# Model Converter Tool
2+
3+
A command-line utility to download PyTorch models and convert them to OpenVINO format.
4+
5+
## Overview
6+
7+
This tool reads a JSON configuration file containing model specifications, downloads PyTorch weights from URLs, loads the models, and exports them to OpenVINO Intermediate Representation (IR) format.
8+
9+
## Features
10+
11+
- **Automatic Download**: Downloads model weights from HTTP/HTTPS URLs with caching support
12+
- **Dynamic Model Loading**: Dynamically imports and instantiates model classes from Python paths
13+
- **Metadata Embedding**: Embeds custom metadata into OpenVINO models
14+
- **Input/Output Naming**: Configurable input and output tensor names
15+
- **Batch Processing**: Process multiple models from a single configuration file
16+
- **Selective Conversion**: Convert specific models using the `--model` flag
17+
18+
## Installation
19+
20+
### Prerequisites
21+
22+
```bash
23+
# Required packages
24+
uv pip install torch torchvision openvino
25+
26+
```
27+
28+
## Usage
29+
30+
### Basic Usage
31+
32+
```bash
33+
uv run python model_converter.py config.json -o ./output_models
34+
```
35+
36+
### Command-Line Options
37+
38+
```text
39+
positional arguments:
40+
config Path to JSON configuration file
41+
42+
options:
43+
-h, --help Show help message and exit
44+
-o OUTPUT, --output OUTPUT
45+
Output directory for converted models (default: ./converted_models)
46+
-c CACHE, --cache CACHE
47+
Cache directory for downloaded weights (default: ~/.cache/torch/hub/checkpoints)
48+
--model MODEL Process only the specified model (by model_short_name)
49+
--list List all models in the configuration file and exit
50+
-v, --verbose Enable verbose logging
51+
```
52+
53+
### Examples
54+
55+
**List all models in configuration:**
56+
57+
```bash
58+
uv run python model_converter.py example_config.json --list
59+
```
60+
61+
**Convert all models:**
62+
63+
```bash
64+
uv run python model_converter.py example_config.json -o ./converted_models
65+
```
66+
67+
**Convert a specific model:**
68+
69+
```bash
70+
uv run python model_converter.py example_config.json -o ./converted_models --model resnet50
71+
```
72+
73+
**Use custom cache directory:**
74+
75+
```bash
76+
uv run python model_converter.py example_config.json -o ./output -c ./my_cache
77+
```
78+
79+
**Enable verbose logging:**
80+
81+
```bash
82+
uv run python model_converter.py example_config.json -o ./output -v
83+
```
84+
85+
## Configuration File Format
86+
87+
The configuration file is a JSON file with the following structure:
88+
89+
```json
90+
{
91+
"models": [
92+
{
93+
"model_short_name": "resnet50",
94+
"model_class_name": "torchvision.models.resnet.resnet50",
95+
"model_full_name": "ResNet-50",
96+
"description": "ResNet-50 image classification model",
97+
"weights_url": "https://download.pytorch.org/models/resnet50-0676ba61.pth",
98+
"input_shape": [1, 3, 224, 224],
99+
"input_names": ["images"],
100+
"output_names": ["output"],
101+
"model_params": null,
102+
"model_type": "Classification"
103+
}
104+
]
105+
}
106+
```
107+
108+
**Important**: The `model_type` field enables automatic model detection when using [Intel's model_api](https://github.com/openvinotoolkit/model_api). When specified, this metadata is embedded in the OpenVINO IR, allowing `Model.create_model()` to automatically select the correct model wrapper class.
109+
110+
Common `model_type` values:
111+
112+
- `"Classification"` - Image classification models
113+
- `"DetectionModel"` - Object detection models
114+
- `"YOLOX"` - YOLOX detection models
115+
- `"SegmentationModel"` - Segmentation models
116+
117+
### Configuration Fields
118+
119+
#### Required Fields
120+
121+
- **`model_short_name`** (string): Short identifier for the model (used for output filename)
122+
- **`model_class_name`** (string): Full Python path to the model class (e.g., `torchvision.models.resnet.resnet50`)
123+
- **`weights_url`** (string): URL to download the PyTorch weights (.pth file)
124+
125+
#### Optional Fields
126+
127+
- **`model_full_name`** (string): Full descriptive name of the model
128+
- **`description`** (string): Description of the model
129+
- **`input_shape`** (array of integers): Input tensor shape (default: `[1, 3, 224, 224]`)
130+
- **`input_names`** (array of strings): Names for input tensors (default: `["input"]`)
131+
- **`output_names`** (array of strings): Names for output tensors (default: auto-generated)
132+
- **`model_params`** (object): Parameters to pass to model constructor (default: `null`)
133+
- **`model_type`** (string): Model type for model_api auto-detection (e.g., `"Classification"`, `"DetectionModel"`, `"YOLOX"`, etc.)

‎tools/model_converter/config.json‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
{
2+
"models": [
3+
{
4+
"model_short_name": "mobilenet_v3_small",
5+
"model_class_name": "torchvision.models.mobilenetv3.mobilenet_v3_small",
6+
"model_full_name": "MobileNetV3-Small",
7+
"description": "MobileNetV3 Small - Efficient convolutional neural network for mobile and embedded vision applications",
8+
"docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.mobilenet_v3_small.html#torchvision.models.mobilenet_v3_small",
9+
"weights_url": "https://download.pytorch.org/models/mobilenet_v3_small-047dcff4.pth",
10+
"input_shape": [1, 3, 224, 224],
11+
"input_names": ["image"],
12+
"output_names": ["output1"],
13+
"model_params": null,
14+
"model_type": "Classification",
15+
"reverse_input_channels": false,
16+
"mean_values": "123.675 116.28 103.53",
17+
"scale_values": "58.395 57.12 57.375",
18+
"labels": "IMAGENET1K_V1"
19+
},
20+
{
21+
"model_short_name": "efficientnet_b0",
22+
"model_class_name": "torchvision.models.efficientnet.efficientnet_b0",
23+
"model_full_name": "EfficientNet-B0",
24+
"description": "EfficientNet-B0 - Efficient convolutional neural network with compound scaling",
25+
"docs": "https://docs.pytorch.org/vision/main/models/generated/torchvision.models.efficientnet_b0.html#torchvision.models.efficientnet_b0",
26+
"weights_url": "https://download.pytorch.org/models/efficientnet_b0_rwightman-3dd342df.pth",
27+
"input_shape": [1, 3, 224, 224],
28+
"input_names": ["image"],
29+
"output_names": ["logits"],
30+
"model_params": null,
31+
"model_type": "Classification",
32+
"reverse_input_channels": true,
33+
"mean_values": "123.675 116.28 103.53",
34+
"scale_values": "58.395 57.12 57.375",
35+
"labels": "IMAGENET1K_V1"
36+
},
37+
{
38+
"model_short_name": "resnet18",
39+
"model_class_name": "torchvision.models.resnet.resnet18",
40+
"model_full_name": "ResNet-18",
41+
"description": "ResNet-18 - 18-layer residual learning network for image classification",
42+
"weights_url": "https://download.pytorch.org/models/resnet18-f37072fd.pth",
43+
"input_shape": [1, 3, 224, 224],
44+
"input_names": ["image"],
45+
"output_names": ["output"],
46+
"model_params": null,
47+
"model_type": "Classification",
48+
"reverse_input_channels": true,
49+
"mean_values": "123.675 116.28 103.53",
50+
"scale_values": "58.395 57.12 57.375",
51+
"labels": "IMAGENET1K_V1"
52+
},
53+
{
54+
"model_short_name": "resnet50",
55+
"model_class_name": "torchvision.models.resnet.resnet50",
56+
"model_full_name": "ResNet-50",
57+
"description": "ResNet-50 - 50-layer residual learning network for image classification",
58+
"weights_url": "https://download.pytorch.org/models/resnet50-0676ba61.pth",
59+
"input_shape": [1, 3, 224, 224],
60+
"input_names": ["image"],
61+
"output_names": ["output"],
62+
"model_params": null,
63+
"model_type": "Classification",
64+
"reverse_input_channels": true,
65+
"mean_values": "123.675 116.28 103.53",
66+
"scale_values": "58.395 57.12 57.375",
67+
"labels": "IMAGENET1K_V1"
68+
},
69+
{
70+
"model_short_name": "squeezenet1_0",
71+
"model_class_name": "torchvision.models.squeezenet.squeezenet1_0",
72+
"model_full_name": "SqueezeNet 1.0",
73+
"description": "SqueezeNet 1.0 - Small CNN with AlexNet-level accuracy and 50x fewer parameters",
74+
"weights_url": "https://download.pytorch.org/models/squeezenet1_0-b66bff10.pth",
75+
"input_shape": [1, 3, 224, 224],
76+
"input_names": ["image"],
77+
"output_names": ["output"],
78+
"model_params": null,
79+
"model_type": "Classification",
80+
"reverse_input_channels": true,
81+
"mean_values": "123.675 116.28 103.53",
82+
"scale_values": "58.395 57.12 57.375",
83+
"labels": "IMAGENET1K_V1"
84+
},
85+
{
86+
"model_short_name": "vgg16",
87+
"model_class_name": "torchvision.models.vgg.vgg16",
88+
"model_full_name": "VGG-16",
89+
"description": "VGG-16 - 16-layer deep convolutional network",
90+
"weights_url": "https://download.pytorch.org/models/vgg16-397923af.pth",
91+
"input_shape": [1, 3, 224, 224],
92+
"input_names": ["image"],
93+
"output_names": ["output"],
94+
"model_params": null,
95+
"model_type": "Classification",
96+
"reverse_input_channels": true,
97+
"mean_values": "123.675 116.28 103.53",
98+
"scale_values": "58.395 57.12 57.375",
99+
"labels": "IMAGENET1K_V1"
100+
}
101+
]
102+
}

0 commit comments

Comments
 (0)