Tutorial for Super-resolution using caffe2

Author: Priya Goyal

_static/img/SRResNet.png

@@ -0,0 +1,200 @@
+"""
+Transfering a model from PyTorch to Caffe2 and Mobile using ONNX
+================================================================
+
+"""
+
+
+######################################################################
+# In this tutorial, we describe how to use ONNX to convert a model defined
+# in PyTorch into the ONNX format and then load it into Caffe2. Once in
+# Caffe2, we can run the model to double-check it was exported correctly,
+# and we then show how to use Caffe2 features such as mobile exporter for
+# executing the model on mobile devices.
+#
+
+# Some standard imports
+import io
+import numpy as np
+
+from torch import nn
+from torch.autograd import Variable
+import torch.utils.model_zoo as model_zoo
+import torch.onnx
+
+import onnx
+import onnx.backend
+import onnx.backend.caffe2
+
+
+######################################################################
+# For this tutorial, we will transfer a super-resolution model as an
+# example. First, let's create a SuperResolution model in PyTorch. `This
+# model <https://github.com/pytorch/examples/blob/master/super_resolution/model.py>`__
+# comes directly from PyTorch's examples without modification:
+#
+
+# Super Resolution model definition in PyTorch
+import torch.nn as nn
+import torch.nn.init as init
+
+
+class SuperResolutionNet(nn.Module):
+    def __init__(self, upscale_factor, inplace=False):
+        super(SuperResolutionNet, self).__init__()
+
+        self.relu = nn.ReLU(inplace=inplace)
+        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
+        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
+        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
+        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
+        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.conv2(x))
+        x = self.relu(self.conv3(x))
+        x = self.pixel_shuffle(self.conv4(x))
+        return x
+
+    def _initialize_weights(self):
+        init.orthogonal(self.conv1.weight, init.calculate_gain('relu'))
+        init.orthogonal(self.conv2.weight, init.calculate_gain('relu'))
+        init.orthogonal(self.conv3.weight, init.calculate_gain('relu'))
+        init.orthogonal(self.conv4.weight)
+
+# Create the super-resolution model by using the above model definition.
+torch_model = SuperResolutionNet(upscale_factor=3)
+
+
+######################################################################
+# Ordinarily, you would now train this model; however, for this tutorial,
+# we will instead download some pre-trained weights.
+#
+
+# Load pretrained model weights
+model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
+batch_size = 1    # just a random number
+
+# Initialize model with the pretrained weights
+torch_model.load_state_dict(model_zoo.load_url(model_url))
+
+# set the train mode to false since we will only run the forward pass.
+torch_model.train(False)
+
+
+######################################################################
+# Exporting a model in PyTorch works via tracing. To export a model, you
+# call the ``torch.onnx._export()`` function. This will execute the model,
+# recording a trace of what operators are used to compute the outputs.
+# Because ``export`` runs the model, we need provide an input tensor
+# ``x``. The values in this tensor are not important; it can be an image
+# or a random tensor as long as it is the right size.
+#
+
+# Input to the model
+x = Variable(torch.randn(batch_size, 1, 224, 224), requires_grad=True)
+
+# Export the model
+torch_out = torch.onnx._export(torch_model,             # model being run
+                               x,                       # model input (or a tuple for multiple inputs)
+                               "super_resolution.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
+
+
+######################################################################
+# ``torch_out`` is the output after executing the model. Normally you can
+# ignore this output, but here we will use it to verify that the model we
+# exported computes the same values when run in Caffe2.
+#
+# Now let's take the ONNX representation and use it in Caffe2. This part
+# can normally be done in a separate process or on another machine, but we
+# will continue in the same process so that we can verify that Caffe2 and
+# PyTorch are computing the same value for the network:
+#
+
+# Load the ONNX GraphProto object. Graph is a standard Python protobuf object
+graph = onnx.load("super_resolution.onnx")
+
+# prepare the caffe2 backend for executing the model this converts the ONNX graph into a
+# Caffe2 NetDef that can execute it. Other ONNX backends, like one for CNTK will be
+# availiable soon.
+prepared_backend = onnx.backend.caffe2.prepare(graph)
+
+# run the model in Caffe2
+
+# Construct a map from input names to Tensor data.
+# The graph itself contains inputs for all weight parameters, followed by the input image.
+# Since the weights are already embedded, we just need to pass the input image.
+# last input the graph
+W = {graph.input[-1]: x.data.numpy()}
+
+# Run the Caffe2 net:
+c2_out = prepared_backend.run(W)[0]
+
+# Verify the numerical correctness upto 3 decimal places
+np.testing.assert_almost_equal(torch_out.data.cpu().numpy(), c2_out, decimal=3)
+
+
+######################################################################
+# Transfering SRResNet using ONNX
+# ===============================
+#
+
+
+######################################################################
+# Super-resolution is a way of increasing the resolution of images, videos
+# and is widely used in image processing or video editing. For the purpose
+# of tutorial, we used a small super-resolution model with a dummy input
+# above but using the steps above, we also transferred the SRResNet model
+# for super-resolution presented in `this
+# paper <https://arxiv.org/pdf/1609.04802.pdf>`__. The model definition
+# and a pre-trained model are available
+# `here <https://gist.github.com/prigoyal/b245776903efbac00ee89699e001c9bd>`__.
+# Below is what SRResNet model input, output looks like. |SRResNet|
+#
+# .. |SRResNet| image:: /_static/img/SRResNet.png
+#
+
+
+######################################################################
+# Running the model on mobile devices
+# ===================================
+#
+
+
+######################################################################
+# So far we have exported a model from PyTorch and shown how to load it
+# and run it in Caffe2. Now that the model is loaded in Caffe2, we can
+# convert it into a format suitable for `running on mobile
+# devices <https://caffe2.ai/docs/mobile-integration.html>`__.
+#
+# We will use Caffe2's
+# `mobile\_exporter <https://github.com/caffe2/caffe2/blob/master/caffe2/python/predictor/mobile_exporter.py>`__
+# to generate the two model protobufs that can run on mobile. The first is
+# used to initialize the network with the correct weights, and the second
+# actual runs executes the model.
+#
+
+# extract the workspace and the graph proto from the internal representation
+c2_workspace = prepared_backend.workspace
+c2_graph = prepared_backend.predict_net
+
+# Now import the caffe2 mobile exporter
+from caffe2.python.predictor import mobile_exporter
+
+# call the Export to get the predict_net, init_net. These nets are needed for running things on mobile
+init_net, predict_net = mobile_exporter.Export(c2_workspace, c2_graph, c2_graph.external_input)
+
+
+######################################################################
+# Now, on your ios/Android device, you can use the above protobufs and use
+# ``caffe2::Predictor`` (iOS) or Caffe2 instance (Android) for deploying
+# them real-time.
+# `Here <https://gist.github.com/prigoyal/6bbec5cd121182596848af1c265d7bbe>`__
+# is the gist showing steps for running the model on mobile and verifying
+# the correctness. Also, for more information, checkout
+# `caffe2-android-demo <https://caffe2.ai/docs/AI-Camera-demo-android.html>`__
+#

advanced_source/super_resolution_with_caffe2.py

@@ -1,27 +1,27 @@
 Welcome to PyTorch Tutorials
 ============================
 
-To get started with learning PyTorch, start with our Beginner Tutorials. 
-The :doc:`60-minute blitz </beginner/deep_learning_60min_blitz>` is the most common 
+To get started with learning PyTorch, start with our Beginner Tutorials.
+The :doc:`60-minute blitz </beginner/deep_learning_60min_blitz>` is the most common
 starting point, and gives you a quick introduction to PyTorch.
-If you like learning by examples, you will like the tutorial 
+If you like learning by examples, you will like the tutorial
 :doc:`/beginner/pytorch_with_examples`
 
-If you would like to do the tutorials interactively via IPython / Jupyter, 
+If you would like to do the tutorials interactively via IPython / Jupyter,
 each tutorial has a download link for a Jupyter Notebook and Python source code.
 
-We also provide a lot of high-quality examples covering image classification, 
-unsupervised learning, reinforcement learning, machine translation and 
+We also provide a lot of high-quality examples covering image classification,
+unsupervised learning, reinforcement learning, machine translation and
 many other applications at https://github.com/pytorch/examples/
 
-You can find reference documentation for PyTorch's API and layers at 
+You can find reference documentation for PyTorch's API and layers at
 http://docs.pytorch.org or via inline help.
 If you would like the tutorials section improved, please open a github issue
 here with your feedback: https://github.com/pytorch/tutorials
 
 Beginner Tutorials
 ------------------
-   
+
 .. customgalleryitem::
    :figure: /_static/img/thumbnails/pytorch-logo-flat.png
    :tooltip: Understand PyTorch’s Tensor library and neural networks at a high level.
@@ -100,7 +100,7 @@ Reinforcement Learning
    :includehidden:
    :hidden:
    :caption: Intermediate Tutorials
-   
+
    intermediate/char_rnn_classification_tutorial
    intermediate/char_rnn_generation_tutorial
    intermediate/seq2seq_translation_tutorial
@@ -132,5 +132,5 @@ Advanced Tutorials
 
    advanced/neural_style_tutorial
    advanced/numpy_extensions_tutorial
+   advanced/super_resolution_with_caffe2
    advanced/c_extension
-