From e68433c46698d09ede9ae4a5792920e8f7eb32bc Mon Sep 17 00:00:00 2001
From: Neta Zmora <31280975+nzmora@users.noreply.github.com>
Date: Sat, 13 Oct 2018 20:55:28 +0300
Subject: [PATCH] ONNX export: add Softmax layer to the end of image
classifiers (#57)
When running inference in ONNX, we often want to add a softmax
layer to TorchVision's models.
---
apputils/model_summaries.py | 9 +++++++++
1 file changed, 9 insertions(+)
diff --git a/apputils/model_summaries.py b/apputils/model_summaries.py
index 2ba1852..3d0cedc 100755
--- a/apputils/model_summaries.py
+++ b/apputils/model_summaries.py
@@ -24,6 +24,7 @@ import os
import re
import numpy as np
import collections
+from copy import deepcopy
import torch
import torchvision
from torch.autograd import Variable
@@ -615,6 +616,14 @@ def export_img_classifier_to_onnx(model, onnx_fname, dataset):
# Pytorch 0.4 doesn't support exporting modules wrapped in DataParallel
if isinstance(model, torch.nn.DataParallel):
model = model.module
+
+ # Explicitly add a softmax layer, because it is needed for the ONNX inference phase.
+ # We make a copy of the model, since we are about to change it (adding softmax).
+ model = deepcopy(model)
+ model.original_forward = model.forward
+ softmax = torch.nn.Softmax(dim=1)
+ model.forward = lambda input: softmax(model.original_forward(input))
+
torch.onnx.export(model, dummy_input, onnx_fname, verbose=False, export_params=True)
msglogger.info('Exported the model to ONNX format at %s' % os.path.realpath(onnx_fname))
--
GitLab