PyTorchでファインチューニングしたモデルをONNXで利用する
昨日の作業の結果、Illustration2Vecのモデルが大きすぎて貧弱なサーバーでは使えないことが分かりました。今のところ二次元画像判別器の特徴量抽出にしか使っていないので、もっと軽いモデルでも代用できるはずです。軽いモデルとして有名なSqueezenetをこれまで集めたデータでファインチューニングして様子を見てみることにします。
ファインチューニングとONNXへのエキスポート
PyTorchのチュートリアルが丁寧に説明してくれているので、これをコピペして継ぎ接ぎするだけです。
- Finetuning Torchvision Models — PyTorch Tutorials 1.2.0 documentation
- (optional) Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime — PyTorch Tutorials 1.9.1+cu102 documentation
継ぎ接ぎしたものがこちらになります。これを実行するとmodel.onnxというファイルが作成されます。
ONNX版Illustration2Vecのモデルサイズが910Mに対して、このモデルは2.8MBなのでだいぶ小さくなりました。精度もだいたい同じくらいだと思います。
from __future__ import print_function from __future__ import division import torch import torch.nn as nn import torch.optim as optim import numpy as np import torchvision from torchvision import datasets, models, transforms import matplotlib.pyplot as plt import time import os import copy print("PyTorch Version: ",torch.__version__) print("Torchvision Version: ",torchvision.__version__) # Top level data directory. Here we assume the format of the directory conforms # to the ImageFolder structure data_dir = "./data/images" # Models to choose from [resnet, alexnet, vgg, squeezenet, densenet, inception] model_name = "squeezenet" # Number of classes in the dataset num_classes = 2 # Batch size for training (change depending on how much memory you have) batch_size = 8 # Number of epochs to train for num_epochs = 1 # Flag for feature extracting. When False, we finetune the whole model, # when True we only update the reshaped layer params feature_extract = True def train_model(model, dataloaders, criterion, optimizer, num_epochs=25, is_inception=False): since = time.time() val_acc_history = [] best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch, num_epochs - 1)) print('-' * 10) # Each epoch has a training and validation phase for phase in ['train', 'val']: if phase == 'train': model.train() # Set model to training mode else: model.eval() # Set model to evaluate mode running_loss = 0.0 running_corrects = 0 # Iterate over data. for inputs, labels in dataloaders[phase]: inputs = inputs.to(device) labels = labels.to(device) # zero the parameter gradients optimizer.zero_grad() # forward # track history if only in train with torch.set_grad_enabled(phase == 'train'): # Get model outputs and calculate loss # Special case for inception because in training it has an auxiliary output. In train # mode we calculate the loss by summing the final output and the auxiliary output # but in testing we only consider the final output. if is_inception and phase == 'train': # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958 outputs, aux_outputs = model(inputs) loss1 = criterion(outputs, labels) loss2 = criterion(aux_outputs, labels) loss = loss1 + 0.4*loss2 else: outputs = model(inputs) loss = criterion(outputs, labels) _, preds = torch.max(outputs, 1) # backward + optimize only if in training phase if phase == 'train': loss.backward() optimizer.step() # statistics running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / len(dataloaders[phase].dataset) epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset) print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc)) # deep copy the model if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = copy.deepcopy(model.state_dict()) if phase == 'val': val_acc_history.append(epoch_acc) print() time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60)) print('Best val Acc: {:4f}'.format(best_acc)) # load best model weights model.load_state_dict(best_model_wts) return model, val_acc_history def set_parameter_requires_grad(model, feature_extracting): if feature_extracting: for param in model.parameters(): param.requires_grad = False def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True): # Initialize these variables which will be set in this if statement. Each of these # variables is model specific. model_ft = None input_size = 0 if model_name == "resnet": """ Resnet18 """ model_ft = models.resnet18(pretrained=use_pretrained) set_parameter_requires_grad(model_ft, feature_extract) num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, num_classes) input_size = 224 elif model_name == "alexnet": """ Alexnet """ model_ft = models.alexnet(pretrained=use_pretrained) set_parameter_requires_grad(model_ft, feature_extract) num_ftrs = model_ft.classifier[6].in_features model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes) input_size = 224 elif model_name == "vgg": """ VGG11_bn """ model_ft = models.vgg11_bn(pretrained=use_pretrained) set_parameter_requires_grad(model_ft, feature_extract) num_ftrs = model_ft.classifier[6].in_features model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes) input_size = 224 elif model_name == "squeezenet": """ Squeezenet """ model_ft = models.squeezenet1_0(pretrained=use_pretrained) set_parameter_requires_grad(model_ft, feature_extract) model_ft.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1,1), stride=(1,1)) model_ft.num_classes = num_classes input_size = 224 elif model_name == "densenet": """ Densenet """ model_ft = models.densenet121(pretrained=use_pretrained) set_parameter_requires_grad(model_ft, feature_extract) num_ftrs = model_ft.classifier.in_features model_ft.classifier = nn.Linear(num_ftrs, num_classes) input_size = 224 elif model_name == "inception": """ Inception v3 Be careful, expects (299,299) sized images and has auxiliary output """ model_ft = models.inception_v3(pretrained=use_pretrained) set_parameter_requires_grad(model_ft, feature_extract) # Handle the auxilary net num_ftrs = model_ft.AuxLogits.fc.in_features model_ft.AuxLogits.fc = nn.Linear(num_ftrs, num_classes) # Handle the primary net num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs,num_classes) input_size = 299 else: print("Invalid model name, exiting...") exit() return model_ft, input_size # Initialize the model for this run model_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True) # Print the model we just instantiated print(model_ft) # Data augmentation and normalization for training # Just normalization for validation data_transforms = { 'train': transforms.Compose([ transforms.RandomResizedCrop(input_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.Resize(input_size), transforms.CenterCrop(input_size), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } print("Initializing Datasets and Dataloaders...") # Create training and validation datasets image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val']} # Create training and validation dataloaders dataloaders_dict = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size, shuffle=True, num_workers=4) for x in ['train', 'val']} # Detect if we have a GPU available device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model_ft = model_ft.to(device) # Gather the parameters to be optimized/updated in this run. If we are # finetuning we will be updating all parameters. However, if we are # doing feature extract method, we will only update the parameters # that we have just initialized, i.e. the parameters with requires_grad # is True. params_to_update = model_ft.parameters() print("Params to learn:") if feature_extract: params_to_update = [] for name,param in model_ft.named_parameters(): if param.requires_grad == True: params_to_update.append(param) print("\t",name) else: for name,param in model_ft.named_parameters(): if param.requires_grad == True: print("\t",name) # Observe that all parameters are being optimized optimizer_ft = optim.SGD(params_to_update, lr=0.001, momentum=0.9) # Setup the loss fxn criterion = nn.CrossEntropyLoss() # Train and evaluate model_ft, hist = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, num_epochs=num_epochs, is_inception=(model_name=="inception")) # Save PyTorch model to file torch.save(model_ft.to('cpu').state_dict(), 'model.pth') # Input to the model x = torch.randn(1, 3, 224, 224, requires_grad=True) torch_out = model_ft(x) # Export the model torch.onnx.export(model_ft, # model being run x, # model input (or a tuple for multiple inputs) "model.onnx", # where to save the model (can be a file or file-like object) export_params=True, # store the trained parameter weights inside the model file opset_version=10, # the ONNX version to export the model to do_constant_folding=True, # whether to execute constant folding for optimization input_names = ['input'], # the model's input names output_names = ['output'], # the model's output names dynamic_axes={'input' : {0 : 'batch_size'}, # variable lenght axes 'output' : {0 : 'batch_size'}})
ONNXモデルの利用
こうして作成したONNXモデルをPyTorchを使わずに利用するコードはこんな感じです。
import os from PIL import Image import numpy as np import onnxruntime # 中心を正方形に切り抜いてリサイズ def crop_and_resize(img, size): width, height = img.size crop_size = min(width, height) img_crop = img.crop(((width - crop_size) // 2, (height - crop_size) // 2, (width + crop_size) // 2, (height + crop_size) // 2)) return img_crop.resize((size, size)) img_mean = np.asarray([0.485, 0.456, 0.406]) img_std = np.asarray([0.229, 0.224, 0.225]) ort_session = onnxruntime.InferenceSession( os.path.join(os.path.dirname(__file__), "model.onnx")) img = Image.open('image.jpg').convert('RGB') img = crop_and_resize(img, 224) # 画像の正規化 img_np = np.asarray(img).astype(np.float32)/255.0 img_np_normalized = (img_np - img_mean) / img_std # (H, W, C) -> (C, H, W) img_np_transposed = img_np_normalized.transpose(2, 0, 1) batch_img = [img_np_transposed] ort_inputs = {ort_session.get_inputs()[0].name: batch_img} ort_outs = ort_session.run(None, ort_inputs)[0] batch_result = np.argmax(ort_outs, axis=1) print(batch_result)
このSqueezenetモデルを使って昨日と同じようなことをするviews.pyが以下のようになります。全文はGitHubを見てください。
github.com
import os import re import urllib.request from urllib.parse import urlparse from PIL import Image from joblib import dump, load import tweepy from django.shortcuts import render from social_django.models import UserSocialAuth from django.conf import settings import more_itertools import numpy as np import onnxruntime import torchvision.transforms as transforms def to_numpy(tensor): return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy() ort_session = onnxruntime.InferenceSession( os.path.join(os.path.dirname(__file__), "model.onnx")) data_transforms = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) def index(request): if request.user.is_authenticated: user = UserSocialAuth.objects.get(user_id=request.user.id) consumer_key = settings.SOCIAL_AUTH_TWITTER_KEY consumer_secret = settings.SOCIAL_AUTH_TWITTER_SECRET access_token = user.extra_data['access_token']['oauth_token'] access_secret = user.extra_data['access_token']['oauth_token_secret'] auth = tweepy.OAuthHandler(consumer_key, consumer_secret) auth.set_access_token(access_token, access_secret) api = tweepy.API(auth) timeline = api.home_timeline(count=200, tweet_mode='extended') tweet_media = [] for tweet in timeline: if 'media' in tweet.entities: tweet_media.append(tweet) batch_size = 4 tweet_illust = [] for batch_tweet in more_itertools.chunked(tweet_media, batch_size): batch_img = [] for tweet in batch_tweet: media_url = tweet.extended_entities['media'][0]['media_url'] filename = os.path.basename(urlparse(media_url).path) filename = os.path.join( os.path.dirname(__file__), 'images', filename) urllib.request.urlretrieve(media_url, filename) img = Image.open(filename).convert('RGB') img = data_transforms(img) batch_img.append(to_numpy(img)) ort_inputs = {ort_session.get_inputs()[0].name: batch_img} ort_outs = ort_session.run(None, ort_inputs)[0] batch_result = np.argmax(ort_outs, axis=1) for tweet, result in zip(batch_tweet, batch_result): if result == 1: media_url = tweet.extended_entities['media'][0]['media_url'] if hasattr(tweet, "retweeted_status"): profile_image_url = tweet.retweeted_status.author.profile_image_url_https author = {'name': tweet.retweeted_status.author.name, 'screen_name': tweet.retweeted_status.author.screen_name} id_str = tweet.retweeted_status.id_str else: profile_image_url = tweet.author.profile_image_url_https author = {'name': tweet.author.name, 'screen_name': tweet.author.screen_name} id_str = tweet.id_str try: text = tweet.retweeted_status.full_text except AttributeError: text = tweet.full_text text = re.sub( r"https?://[\w/:%#\$&\?\(\)~\.=\+\-]+$", '', text).rstrip() tweet_illust.append({'id_str': id_str, 'profile_image_url': profile_image_url, 'author': author, 'text': text, 'image_url': media_url}) tweet_illust_chunked = list(more_itertools.chunked(tweet_illust, 4)) return render(request, 'hello/index.html', {'user': user, 'timeline_chunked': tweet_illust_chunked}) else: return render(request, 'hello/index.html')
モデルがだいぶ小さくなったので、貧弱サーバーでも動かすことができました。
デプロイ
このモデルサイズならHerokuで動かせると思ったのですが、torchvisionなどの依存ライブラリの容量だけでHerokuの500MBの制限を超えてしまうようなので自分のサーバーで動かすことにしました。一応動いてはいるのですが、コールバックの設定がうまくいかないので後日直します。
2.9MBのモデルなら一瞬で推論できると思ったのですが、それでもまだ貧弱サーバーには荷が重いようで読み込みにだいぶ時間がかかります。もっと軽いモデルを作るかディープラーニングに頼らない方法を考えるのが良さそうです。
PyTorchのCPU版をrequirements.txtで指定すればHerokuにデプロイできました。
追記: ヒストグラムの表示
最適なしきい値を見つけるためのヒストグラム表示は次のようにすればできます。
import os import matplotlib.pyplot as plt from PIL import Image import numpy as np import onnxruntime img_mean = np.asarray([0.485, 0.456, 0.406]) img_std = np.asarray([0.229, 0.224, 0.225]) ort_session = onnxruntime.InferenceSession( os.path.join(os.path.dirname(__file__), "model.onnx")) # https://stackoverflow.com/questions/34968722/ def softmax(x): """Compute softmax values for each sets of scores in x.""" e_x = np.exp(x - np.max(x)) return e_x / e_x.sum(axis=0) probs = [] path = 'test' files = os.listdir(path) filenames = [os.path.join(path, f) for f in files if os.path.isfile(os.path.join(path, f))] for name in filenames: img = Image.open(name).convert('RGB') img = img.resize((224, 224)) # Normalization img_np = np.asarray(img).astype(np.float32)/255.0 img_np_normalized = (img_np - img_mean) / img_std # (H, W, C) -> (C, H, W) img_np_transposed = img_np_normalized.transpose(2, 0, 1) batch_img = [img_np_transposed] ort_inputs = {ort_session.get_inputs()[0].name: batch_img} ort_outs = ort_session.run(None, ort_inputs)[0][0] prob = softmax(ort_outs)[1] probs.append(prob) plt.hist(probs) plt.show()