Skip to content

Commit 8e51a7e

Browse files
committed
tools/model_converter
1 parent da6f46e commit 8e51a7e

File tree

3 files changed

+804
-0
lines changed

3 files changed

+804
-0
lines changed

tools/model_converter/README.md

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
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+
python model_converter.py config.json -o ./output_models
34+
```
35+
36+
### Command-Line Options
37+
38+
```
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+
```bash
57+
python model_converter.py example_config.json --list
58+
```
59+
60+
**Convert all models:**
61+
```bash
62+
python model_converter.py example_config.json -o ./converted_models
63+
```
64+
65+
**Convert a specific model:**
66+
```bash
67+
python model_converter.py example_config.json -o ./converted_models --model resnet50
68+
```
69+
70+
**Use custom cache directory:**
71+
```bash
72+
python model_converter.py example_config.json -o ./output -c ./my_cache
73+
```
74+
75+
**Enable verbose logging:**
76+
```bash
77+
python model_converter.py example_config.json -o ./output -v
78+
```
79+
80+
## Configuration File Format
81+
82+
The configuration file is a JSON file with the following structure:
83+
84+
```json
85+
{
86+
"models": [
87+
{
88+
"model_short_name": "resnet50",
89+
"model_class_name": "torchvision.models.resnet.resnet50",
90+
"model_full_name": "ResNet-50",
91+
"description": "ResNet-50 image classification model",
92+
"weights_url": "https://download.pytorch.org/models/resnet50-0676ba61.pth",
93+
"input_shape": [1, 3, 224, 224],
94+
"input_names": ["images"],
95+
"output_names": ["output"],
96+
"model_params": null,
97+
"model_type": "Classification"
98+
}
99+
]
100+
}
101+
```
102+
103+
**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.
104+
105+
Common `model_type` values:
106+
- `"Classification"` - Image classification models
107+
- `"DetectionModel"` - Object detection models
108+
- `"YOLOX"` - YOLOX detection models
109+
- `"SegmentationModel"` - Segmentation models
110+
111+
### Configuration Fields
112+
113+
#### Required Fields
114+
115+
- **`model_short_name`** (string): Short identifier for the model (used for output filename)
116+
- **`model_class_name`** (string): Full Python path to the model class (e.g., `torchvision.models.resnet.resnet50`)
117+
- **`weights_url`** (string): URL to download the PyTorch weights (.pth file)
118+
119+
#### Optional Fields
120+
121+
- **`model_full_name`** (string): Full descriptive name of the model
122+
- **`description`** (string): Description of the model
123+
- **`input_shape`** (array of integers): Input tensor shape (default: `[1, 3, 224, 224]`)
124+
- **`input_names`** (array of strings): Names for input tensors (default: `["input"]`)
125+
- **`output_names`** (array of strings): Names for output tensors (default: auto-generated)
126+
- **`model_params`** (object): Parameters to pass to model constructor (default: `null`)
127+
- **`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)