NetsPresso tutorial for PyTorch_CIFAR_Models compression

Order of the tutorial

0. Sign up
1. Install
2. Prepare the dataset
3. (Optioanl) Run tensorboard
4. Training
5. Testing
6. Model compression with NetsPresso Python Package
7. Fine-tuning the compressed model

0. Sign up

To get started with the NetsPresso Python package, you will need to sign up at NetsPresso.

1. Install

Clone repo, including PyTorch >=1.11, < 2.0.

git clone https://github.com/Nota-NetsPresso/ModelZoo-PyTorch-CIFAR100  # clone
cd ModelZoo-PyTorch-CIFAR100

2. Prepare the dataset

We will use cifar100 dataset from torchvision since it's more convenient, but we also kept the sample code for writing your dataset module in dataset folder, as an example for people don't know how to write it.

3. (Optional) Run tensorbard

Install tensorboard

pip install tensorboard
mkdir runs
Run tensorboard
tensorboard --logdir='runs' --port=6006 --host='localhost'

4. Training

You need to specify the net you want to train using arg -net

# use gpu to train vgg16
python train.py -net vgg16 -gpu

Sometimes, you might want to use warmup training by set -warm to 1 or 2, to prevent network diverge during early training phase.

The supported net args are:

vgg16
repvgg
mobilenetv2
resnet56
inceptionv3

or
saved model path

The pretrained models are from here.

5. Testing

Test the model using test.py

python test.py -net path_to_fx_model_file

6. Model compression with NetsPresso Python Package

Upload & compress your 'model-epoch-fx.pt' by using NetsPresso Python Package

6_1. Install NetsPresso Python Package

pip install netspresso

6_2. Upload & Compress

First, import the packages and set a NetsPresso username and password.

from netspresso.compressor import ModelCompressor, Task, Framework, CompressionMethod, RecommendationMethod


EMAIL = "YOUR_EMAIL"
PASSWORD = "YOUR_PASSWORD"
compressor = ModelCompressor(email=EMAIL, password=PASSWORD)

Second, upload 'model-epoch-fx.pt', which is the model converted to torchfx in step 4, with the following code.

# Upload Model
UPLOAD_MODEL_NAME = "PyTorch_CIFAR_model"
TASK = Task.IMAGE_CLASSIFICATION
FRAMEWORK = Framework.PYTORCH
UPLOAD_MODEL_PATH = "./model-epoch-fx.pt"
INPUT_SHAPES = [{"batch": 1, "channel": 3, "dimension": [32, 32]}]
model = compressor.upload_model(
    model_name=UPLOAD_MODEL_NAME,
    task=TASK,
    framework=FRAMEWORK,
    file_path=UPLOAD_MODEL_PATH,
    input_shapes=INPUT_SHAPES,
)

Finally, you can compress the uploaded model with the desired options through the following code.

# Recommendation Compression
COMPRESSED_MODEL_NAME = "test_l2norm"
COMPRESSION_METHOD = CompressionMethod.PR_L2
RECOMMENDATION_METHOD = RecommendationMethod.SLAMP
RECOMMENDATION_RATIO = 0.6
OUTPUT_PATH = "./compressed_model.pt"
compressed_model = compressor.recommendation_compression(
    model_id=model.model_id,
    model_name=COMPRESSED_MODEL_NAME,
    compression_method=COMPRESSION_METHOD,
    recommendation_method=RECOMMENDATION_METHOD,
    recommendation_ratio=RECOMMENDATION_RATIO,
    output_path=OUTPUT_PATH,
)

Click to check 'Full upload & compress code'

pip install netspresso

from netspresso.compressor import ModelCompressor, Task, Framework, CompressionMethod, RecommendationMethod


EMAIL = "YOUR_EMAIL"
PASSWORD = "YOUR_PASSWORD"
compressor = ModelCompressor(email=EMAIL, password=PASSWORD)

# Upload Model
UPLOAD_MODEL_NAME = "PyTorch_CIFAR_model"
TASK = Task.IMAGE_CLASSIFICATION
FRAMEWORK = Framework.PYTORCH
UPLOAD_MODEL_PATH = "./model-epoch-fx.pt"
INPUT_SHAPES = [{"batch": 1, "channel": 3, "dimension": [32, 32]}]
model = compressor.upload_model(
    model_name=UPLOAD_MODEL_NAME,
    task=TASK,
    framework=FRAMEWORK,
    file_path=UPLOAD_MODEL_PATH,
    input_shapes=INPUT_SHAPES,
)

# Recommendation Compression
COMPRESSED_MODEL_NAME = "test_l2norm"
COMPRESSION_METHOD = CompressionMethod.PR_L2
RECOMMENDATION_METHOD = RecommendationMethod.SLAMP
RECOMMENDATION_RATIO = 0.6
OUTPUT_PATH = "./compressed_model.pt"
compressed_model = compressor.recommendation_compression(
    model_id=model.model_id,
    model_name=COMPRESSED_MODEL_NAME,
    compression_method=COMPRESSION_METHOD,
    recommendation_method=RECOMMENDATION_METHOD,
    recommendation_ratio=RECOMMENDATION_RATIO,
    output_path=OUTPUT_PATH,
)

More commands can be found in the official NetsPresso Python Package docs: https://nota-netspresso.github.io/PyNetsPresso-docs

Alternatively, you can do the same as above through the GUI on our website: https://console.netspresso.ai/models

7. Fine-tuning the compressed model

You may need to set learning rate using -lr

$ python train.py -net path_to_compressed_model_file -lr 0.001

Now you can use the compressed model however you like!

Contact

Join our Discussion Forum for providing feedback or sharing your use cases, and if you want to talk more with Nota, please contact us here.
Or you can also do it via email([email protected]) or phone(+82 2-555-8659)!