<jupyter_start><jupyter_text><jupyter_code>!pip install git+https://github.com/alibaba/TinyNeuralNetwork.git
!wget --no-check-certificate \
    https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip \
    -O /content/cats_and_dogs_filtered.zip
import os
import zipfile

local_zip = '/content/cats_and_dogs_filtered.zip'
zip_ref = zipfile.ZipFile(local_zip, 'r')
zip_ref.extractall('/content')
zip_ref.close()
import random

from glob import glob

from PIL import Image
import torch
from torchvision import transforms
import torchvision.models as models


from tinynn.converter import TFLiteConverter
from tinynn.graph.quantization.quantizer import PostQuantizer
from tinynn.graph.tracer import model_tracer
from tinynn.util.cifar10 import get_dataloader, train_one_epoch, validate
from tinynn.util.train_util import DLContext, get_device, train


random.seed(0)


with model_tracer():
  model = models.mobilenet_v2(pretrained=True)
  model.eval()

  # Provide a viable input for the model
  dummy_input = torch.rand((1, 3, 224, 224))

  quantizer = PostQuantizer(model, dummy_input, work_dir='out', config={'asymmetric': True, 'per_tensor': False})
  qat_model = quantizer.quantize()

print(qat_model)

# Use DataParallel to speed up training when possible
if torch.cuda.device_count() > 1:
  qat_model = nn.DataParallel(qat_model)

# Move model to the appropriate device
device = get_device()
qat_model.to(device=device)


dataset_list = glob('/content/cats_and_dogs_filtered/train/**/*', recursive=True)
random.shuffle(dataset_list)
for i in range(100):
  filename = dataset_list[i]      
  print(filename)
  input_image = Image.open(filename)
  preprocess = transforms.Compose([
      transforms.Resize(256),
      transforms.CenterCrop(224),
      transforms.ToTensor(),
      transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
  ])
  input_tensor = preprocess(input_image)
  print(input_tensor.shape)
  input_tensor = torch.unsqueeze(input_tensor, 0)
  print("torch input_tensor size")
  print(input_tensor.shape)    
  qat_model(input_tensor.to(device=device))
  

with torch.no_grad():
  qat_model.eval()
  qat_model.cpu()

  # The step below converts the model to an actual quantized model, which uses the quantized kernels.
  qat_model = torch.quantization.convert(qat_model)

  # When converting quantized models, please ensure the quantization backend is set.
  torch.backends.quantized.engine = quantizer.backend

  # The code section below is used to convert the model to the TFLite format
  # If you need a quantized model with a specific data type (e.g. int8)
  # you may specify `quantize_target_type='int8'` in the following line.
  # If you need a quantized model with strict symmetric quantization check (with pre-defined zero points),
  # you may specify `strict_symmetric_check=True` in the following line.
  converter = TFLiteConverter(qat_model, dummy_input, tflite_path='out/qat_model.tflite', quantize_target_type='int8', input_transpose=False, fuse_quant_dequant=True)
  converter.convert()


# Download an example image from the pytorch website
import urllib
url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
try: urllib.URLopener().retrieve(url, filename)
except: urllib.request.urlretrieve(url, filename)

import tensorflow as tf
import numpy as np
tflite_model_path = '/content/out/qat_model.tflite'
interpreter = tf.lite.Interpreter(model_path=tflite_model_path)
interpreter.allocate_tensors()

# Get input and output tensors
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
test_details = interpreter.get_input_details()[0]

scale, zero_point = test_details['quantization']
print(scale)
print(zero_point)

# Test the model on image  data
# sample execution (requires torchvision)
from PIL import Image
from torchvision import transforms
input_image = Image.open(filename)
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
print(input_tensor.shape)
input_tensor = torch.unsqueeze(input_tensor, 0)
input_tensor = torch.quantize_per_tensor(input_tensor, torch.tensor(scale), torch.tensor(zero_point), torch.qint8)
input_tensor = torch.int_repr(input_tensor).numpy()

print("torch input_tensor size:")
print(input_tensor.shape)
print(input_tensor)
interpreter.set_tensor(input_details[0]['index'], input_tensor)

interpreter.invoke()

# get_tensor() returns a copy of the tensor data
# use tensor() in order to get a pointer to the tensor
output_data = interpreter.get_tensor(output_details[0]['index'])

print("Predicted value . Label index: {}, confidence: {:2.0f}%"
      .format(np.argmax(output_data), 
              100 * output_data[0][np.argmax(output_data)]))<jupyter_output>0.01871182955801487
-15
torch.Size([3, 224, 224])
torch input_tensor size:
(1, 3, 224, 224)
[[[[-118 -118 -117 ... -124 -119 -118]
   [-122 -116 -117 ... -120 -118 -110]
   [-122 -119 -117 ... -125 -120 -116]
   ...
   [ -94 -101 -102 ...  -61  -74  -72]
   [ -97 -101 -102 ...  -71 -102  -92]
   [ -98  -94  -82 ...  -64  -83  -82]]

  [[-113 -113 -112 ... -121 -117 -117]
   [-113 -114 -114 ... -120 -118 -115]
   [-113 -114 -114 ... -120 -119 -117]
   ...
   [ -68  -67  -68 ...  -40  -49  -51]
   [ -68  -68  -69 ...  -45  -75  -63]
   [ -66  -68  -58 ...  -36  -58  -55]]

  [[-101 -100  -98 ... -107 -106 -107]
   [-104 -102 -102 ... -107 -107 -103]
   [-103 -104 -103 ... -108 -106 -105]
   ...
   [ -85  -96 -100 ...  -63  -77  -73]
   [ -84  -97 -101 ...  -66  -95  -89]
   [ -83  -91  -81 ...  -55  -75  -83]]]]
Predicted value . Label index: 258, confidence: 7400%
{'fraction_non_alphanumeric': 0.13196834136269786, 'fraction_numerical': 0.07037164487267722, 'mean_word_length': 3.893097643097643, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 5, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '588319', 'license': 'Apache-2.0', 'n_tokens_mistral': 2309, 'n_tokens_neox': 2016, 'n_words': 507, 'path': 'third_party/xtensa/examples/pytorch_to_tflite/pytorch_to_tflite_converter/tinynn_pytorch_to_tflite_int8.ipynb', 'repo_name': 'chris-kuo/tflite-micro', 'size': '5812'}