How to make vgg pytorch's ptl size smaller on android?

Viewed 32
import torch
# import joblib
from torch.utils.mobile_optimizer import optimize_for_mobile
from torchvision.models.vgg import vgg16
import torch, torchvision.models

# lb = joblib.load('lb.pkl')
device = torch.device('cuda:0')
#device = torch.device('cpu')#'cuda:0')
torch.backends.cudnn.benchmark = True
model = vgg16().to(device)



# model = torchvision.models.vgg16()
path = 'model-22222.pth'
torch.save(model.state_dict(), path) # nothing else here
model.load_state_dict(torch.load(path))

#model.load_state_dict(torch.load('./model-76-0.7754.pth'))

scripted_module = torch.jit.script(model)
optimized_scripted_module = optimize_for_mobile(scripted_module)
optimized_scripted_module._save_for_lite_interpreter("model-76-0.7754.ptl")

The optimize_for_mobile seems does not make the ptl file smaller, it's about 527M, too large on Android. How to make it smaller?

1 Answers

You may try using quantization: https://pytorch.org/docs/stable/quantization.html

Usually it allows to reduce the size of the model 2x-4x times with little or no loss of accuracy.

Example:

import torch
from torch.utils.mobile_optimizer import optimize_for_mobile
from torchvision.models.vgg import vgg16

device = torch.device("cpu")
torch.backends.cudnn.benchmark = True
model = vgg16().to(device)

backend = "qnnpack"
model.qconfig = torch.quantization.get_default_qconfig(backend)
torch.backends.quantized.engine = backend
model_static_quantized = torch.quantization.prepare(model, inplace=False)
model_static_quantized = torch.quantization.convert(
    model_static_quantized, inplace=False
)

scripted_module = torch.jit.script(model_static_quantized)
optimized_scripted_module = optimize_for_mobile(scripted_module)
optimized_scripted_module._save_for_lite_interpreter("model_quantized.ptl")

Apart from it I would try to use the architectures more suitable for mobile phones because VGG is kinda old and has not a great size/accuracy ratio. You can distillate your trained VGG model into smaller one with Knowledge distillation.

Related