Skip to content

Commit df1b6be

Browse files
winnietsanghouseroad
winnietsang
authored and
houseroad
committed
Add a PyTorch to Tensorflow tutorial (onnx#68)
* Add a Pytorch to Tensorflow end-to-end tutorial * Add reference of the new tutorial in README.md file * Remove the return data in the notebook
1 parent 456afe7 commit df1b6be

File tree

4 files changed

+339
-0
lines changed

4 files changed

+339
-0
lines changed

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
* [Serving PyTorch Models on AWS Lambda with Caffe2 & ONNX](https://machinelearnings.co/serving-pytorch-models-on-aws-lambda-with-caffe2-onnx-7b096806cfac)
2424
* [Serving ONNX models with MXNet Model Server](tutorials/ONNXMXNetServer.ipynb)
2525
* [Converting Style Transfer model from PyTorch to CoreML and deploying to an iPhone](https://github.com/onnx/tutorials/tree/master/examples/CoreML/ONNXLive)
26+
* [Convert a PyTorch model to Tensorflow using ONNX](tutorials/PytorchTensorflowMnist.ipynb)
2627

2728
## ONNX tools
2829

‎tutorials/PytorchTensorflowMnist.ipynb‎

Lines changed: 338 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,338 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"metadata": {},
6+
"source": [
7+
"# Convert a PyTorch model to Tensorflow using ONNX\n",
8+
"\n",
9+
"In this tutorial, we will show you how to export a model defined in PyTorch to ONNX and then import the ONNX model into Tensorflow to run it. We will also show you how to save this Tensorflow model into a file for later use. "
10+
]
11+
},
12+
{
13+
"cell_type": "markdown",
14+
"metadata": {},
15+
"source": [
16+
"## Installations\n",
17+
"\n",
18+
"First let's install [ONNX](https://github.com/onnx/onnx), [PyTorch](https://github.com/pytorch/pytorch), and [Tensorflow](https://github.com/tensorflow/tensorflow) by following the instructions on each of their repository.\n",
19+
"\n",
20+
"Then Install torchvision by the following command:\n",
21+
"```\n",
22+
"pip install torchvision\n",
23+
"```\n",
24+
"\n",
25+
"Next install [onnx-tensorflow](https://github.com/onnx/onnx-tensorflow) by the following commands:\n",
26+
"```\n",
27+
"git clone git@github.com:onnx/onnx-tensorflow.git && cd onnx-tensorflow\n",
28+
"pip install -e .\n",
29+
"```"
30+
]
31+
},
32+
{
33+
"cell_type": "markdown",
34+
"metadata": {},
35+
"source": [
36+
"## Define model\n",
37+
"\n",
38+
"In this tutorial we are going to use the [MNIST model](https://github.com/pytorch/examples/tree/master/mnist) from PyTorch examples. "
39+
]
40+
},
41+
{
42+
"cell_type": "code",
43+
"execution_count": null,
44+
"metadata": {},
45+
"outputs": [],
46+
"source": [
47+
"import torch\n",
48+
"import torch.nn as nn\n",
49+
"import torch.nn.functional as F\n",
50+
"\n",
51+
"\n",
52+
"class Net(nn.Module):\n",
53+
" def __init__(self):\n",
54+
" super(Net, self).__init__()\n",
55+
" self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
56+
" self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
57+
" self.conv2_drop = nn.Dropout2d()\n",
58+
" self.fc1 = nn.Linear(320, 50)\n",
59+
" self.fc2 = nn.Linear(50, 10)\n",
60+
"\n",
61+
" def forward(self, x):\n",
62+
" x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
63+
" x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
64+
" x = x.view(-1, 320)\n",
65+
" x = F.relu(self.fc1(x))\n",
66+
" x = F.dropout(x, training=self.training)\n",
67+
" x = self.fc2(x)\n",
68+
" return F.log_softmax(x, dim=1)"
69+
]
70+
},
71+
{
72+
"cell_type": "markdown",
73+
"metadata": {},
74+
"source": [
75+
"## Train and test model\n",
76+
"Now let's train this model. By default if GPU is availalbe on your environment, it will use GPU instead of CPU to run the training. In this tutorial we will train this model with 60 epochs. It will takes about 15 minutes on an environment with 1 GPU to complete this training. You can always adjust the number of epoch base on how well you want your model to be trained."
77+
]
78+
},
79+
{
80+
"cell_type": "code",
81+
"execution_count": null,
82+
"metadata": {},
83+
"outputs": [],
84+
"source": [
85+
"import argparse\n",
86+
"import torch.optim as optim\n",
87+
"from torchvision import datasets, transforms\n",
88+
"\n",
89+
"def train(args, model, device, train_loader, optimizer, epoch):\n",
90+
" model.train()\n",
91+
" for batch_idx, (data, target) in enumerate(train_loader):\n",
92+
" data, target = data.to(device), target.to(device)\n",
93+
" optimizer.zero_grad()\n",
94+
" output = model(data)\n",
95+
" loss = F.nll_loss(output, target)\n",
96+
" loss.backward()\n",
97+
" optimizer.step()\n",
98+
" if batch_idx % args.log_interval == 0:\n",
99+
" print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
100+
" epoch, batch_idx * len(data), len(train_loader.dataset),\n",
101+
" 100. * batch_idx / len(train_loader), loss.item()))\n",
102+
"\n",
103+
"def test(args, model, device, test_loader):\n",
104+
" model.eval()\n",
105+
" test_loss = 0\n",
106+
" correct = 0\n",
107+
" with torch.no_grad():\n",
108+
" for data, target in test_loader:\n",
109+
" data, target = data.to(device), target.to(device)\n",
110+
" output = model(data)\n",
111+
" test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss\n",
112+
" pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability\n",
113+
" correct += pred.eq(target.view_as(pred)).sum().item()\n",
114+
"\n",
115+
" test_loss /= len(test_loader.dataset)\n",
116+
" print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n",
117+
" test_loss, correct, len(test_loader.dataset),\n",
118+
" 100. * correct / len(test_loader.dataset)))\n",
119+
" \n",
120+
"# Training settings\n",
121+
"parser = argparse.ArgumentParser(description='PyTorch MNIST Example')\n",
122+
"parser.add_argument('--batch-size', type=int, default=64, metavar='N',\n",
123+
" help='input batch size for training (default: 64)')\n",
124+
"parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',\n",
125+
" help='input batch size for testing (default: 1000)')\n",
126+
"parser.add_argument('--epochs', type=int, default=10, metavar='N',\n",
127+
" help='number of epochs to train (default: 10)')\n",
128+
"parser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n",
129+
" help='learning rate (default: 0.01)')\n",
130+
"parser.add_argument('--momentum', type=float, default=0.5, metavar='M',\n",
131+
" help='SGD momentum (default: 0.5)')\n",
132+
"parser.add_argument('--no-cuda', action='store_true', default=False,\n",
133+
" help='disables CUDA training')\n",
134+
"parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
135+
" help='random seed (default: 1)')\n",
136+
"parser.add_argument('--log-interval', type=int, default=10, metavar='N',\n",
137+
" help='how many batches to wait before logging training status')\n",
138+
"\n",
139+
"# Train this model with 60 epochs and after process every 300 batches log the train status \n",
140+
"args = parser.parse_args(['--epochs', '60', '--log-interval', '300'])\n",
141+
"\n",
142+
"use_cuda = not args.no_cuda and torch.cuda.is_available()\n",
143+
"\n",
144+
"torch.manual_seed(args.seed)\n",
145+
"\n",
146+
"device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
147+
"\n",
148+
"kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}\n",
149+
"train_loader = torch.utils.data.DataLoader(\n",
150+
" datasets.MNIST('../data', train=True, download=True,\n",
151+
" transform=transforms.Compose([\n",
152+
" transforms.ToTensor(),\n",
153+
" transforms.Normalize((0.1307,), (0.3081,))\n",
154+
" ])),\n",
155+
" batch_size=args.batch_size, shuffle=True, **kwargs)\n",
156+
"test_loader = torch.utils.data.DataLoader(\n",
157+
" datasets.MNIST('../data', train=False, transform=transforms.Compose([\n",
158+
" transforms.ToTensor(),\n",
159+
" transforms.Normalize((0.1307,), (0.3081,))\n",
160+
" ])),\n",
161+
" batch_size=args.test_batch_size, shuffle=True, **kwargs)\n",
162+
"\n",
163+
"\n",
164+
"model = Net().to(device)\n",
165+
"optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)\n",
166+
"\n",
167+
"for epoch in range(1, args.epochs + 1):\n",
168+
" train(args, model, device, train_loader, optimizer, epoch)\n",
169+
" test(args, model, device, test_loader)"
170+
]
171+
},
172+
{
173+
"cell_type": "markdown",
174+
"metadata": {},
175+
"source": [
176+
"## Save the trained model"
177+
]
178+
},
179+
{
180+
"cell_type": "code",
181+
"execution_count": null,
182+
"metadata": {},
183+
"outputs": [],
184+
"source": [
185+
"# Save the trained model to a file\n",
186+
"torch.save(model.state_dict(), 'output/mnist.pth')"
187+
]
188+
},
189+
{
190+
"cell_type": "markdown",
191+
"metadata": {},
192+
"source": [
193+
"## Export the trained model to ONNX \n",
194+
"In order to export the model, Pytorch exporter needs to run the model once and save this resulting traced model to a ONNX file. Therefore, we need to provide the input for our MNIST model. Here we would expect to get a black and white 28 x 28 picture as an input to run this model in inference phase."
195+
]
196+
},
197+
{
198+
"cell_type": "code",
199+
"execution_count": null,
200+
"metadata": {},
201+
"outputs": [],
202+
"source": [
203+
"from torch.autograd import Variable\n",
204+
"\n",
205+
"# Load the trained model from file\n",
206+
"trained_model = Net()\n",
207+
"trained_model.load_state_dict(torch.load('output/mnist.pth'))\n",
208+
"\n",
209+
"# Export the trained model to ONNX\n",
210+
"dummy_input = Variable(torch.randn(1, 1, 28, 28)) # one black and white 28 x 28 picture will be the input to the model\n",
211+
"torch.onnx.export(trained_model, dummy_input, \"output/mnist.onnx\")"
212+
]
213+
},
214+
{
215+
"cell_type": "markdown",
216+
"metadata": {},
217+
"source": [
218+
"PS. You can examine the graph of this mnist.onnx file using an ONNX viewer call [Netron](https://github.com/lutzroeder/Netron)"
219+
]
220+
},
221+
{
222+
"cell_type": "markdown",
223+
"metadata": {},
224+
"source": [
225+
"## Import the ONNX model to Tensorflow\n",
226+
"We will use onnx_tf.backend.prepare to import the ONNX model into Tensorflow."
227+
]
228+
},
229+
{
230+
"cell_type": "code",
231+
"execution_count": null,
232+
"metadata": {},
233+
"outputs": [],
234+
"source": [
235+
"import onnx\n",
236+
"from onnx_tf.backend import prepare\n",
237+
"\n",
238+
"# Load the ONNX file\n",
239+
"model = onnx.load('output/mnist.onnx')\n",
240+
"\n",
241+
"# Import the ONNX model to Tensorflow\n",
242+
"tf_rep = prepare(model)"
243+
]
244+
},
245+
{
246+
"cell_type": "markdown",
247+
"metadata": {},
248+
"source": [
249+
"Let's explore the tf_rep object return from onnx.tf.backend.prepare"
250+
]
251+
},
252+
{
253+
"cell_type": "code",
254+
"execution_count": null,
255+
"metadata": {},
256+
"outputs": [],
257+
"source": [
258+
"# Input nodes to the model\n",
259+
"print('inputs:', tf_rep.inputs)\n",
260+
"\n",
261+
"# Output nodes from the model\n",
262+
"print('outputs:', tf_rep.outputs)\n",
263+
"\n",
264+
"# All nodes in the model\n",
265+
"print('tensor_dict:')\n",
266+
"print(tf_rep.tensor_dict)"
267+
]
268+
},
269+
{
270+
"cell_type": "markdown",
271+
"metadata": {},
272+
"source": [
273+
"## Run the model in Tensorflow"
274+
]
275+
},
276+
{
277+
"cell_type": "code",
278+
"execution_count": null,
279+
"metadata": {},
280+
"outputs": [],
281+
"source": [
282+
"import numpy as np\n",
283+
"from IPython.display import display\n",
284+
"from PIL import Image\n",
285+
"\n",
286+
"print('Image 1:')\n",
287+
"img = Image.open('assets/two.png').resize((28, 28)).convert('L')\n",
288+
"display(img)\n",
289+
"output = tf_rep.run(np.asarray(img, dtype=np.float32)[np.newaxis, np.newaxis, :, :])\n",
290+
"print('The digit is classified as ', np.argmax(output))\n",
291+
"\n",
292+
"print('Image 2:')\n",
293+
"img = Image.open('assets/three.png').resize((28, 28)).convert('L')\n",
294+
"display(img)\n",
295+
"output = tf_rep.run(np.asarray(img, dtype=np.float32)[np.newaxis, np.newaxis, :, :])\n",
296+
"print('The digit is classified as ', np.argmax(output))"
297+
]
298+
},
299+
{
300+
"cell_type": "markdown",
301+
"metadata": {},
302+
"source": [
303+
"## Save the Tensorflow model into a file"
304+
]
305+
},
306+
{
307+
"cell_type": "code",
308+
"execution_count": null,
309+
"metadata": {},
310+
"outputs": [],
311+
"source": [
312+
"tf_rep.export_graph('output/mnist.pb')"
313+
]
314+
}
315+
],
316+
"metadata": {
317+
"celltoolbar": "Raw Cell Format",
318+
"kernelspec": {
319+
"display_name": "Python 2",
320+
"language": "python",
321+
"name": "python2"
322+
},
323+
"language_info": {
324+
"codemirror_mode": {
325+
"name": "ipython",
326+
"version": 3
327+
},
328+
"file_extension": ".py",
329+
"mimetype": "text/x-python",
330+
"name": "python",
331+
"nbconvert_exporter": "python",
332+
"pygments_lexer": "ipython3",
333+
"version": "3.5.2"
334+
}
335+
},
336+
"nbformat": 4,
337+
"nbformat_minor": 2
338+
}

‎tutorials/assets/three.png‎

1.33 KB
Loading

‎tutorials/assets/two.png‎

1.53 KB
Loading

0 commit comments

Comments
 (0)