PyTorch Environment

This notebook builds a reusable environment for using the PyTorch machine learning ecosystem, based on the default Python 3 environment. Nextjournal's PyTorch environment runs PyTorch v1.2.0, and is configured to use Nvidia CUDA v9.2.

Learn more about environments on Nextjournal.

You can use this environment by remixing the Nextjournal PyTorch template.



Adapted from the PyTorch MNIST example. Imports:

from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

Define network.

class Net(nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    self.conv1 = nn.Conv2d(1, 20, 5, 1)
    self.conv2 = nn.Conv2d(20, 50, 5, 1)
    self.fc1 = nn.Linear(4*4*50, 500)
    self.fc2 = nn.Linear(500, 10)

  def forward(self, x):
    x = F.relu(self.conv1(x))
    x = F.max_pool2d(x, 2, 2)
    x = F.relu(self.conv2(x))
    x = F.max_pool2d(x, 2, 2)
    x = x.view(-1, 4*4*50)
    x = F.relu(self.fc1(x))
    x = self.fc2(x)
    return F.log_softmax(x, dim=1)

Train and test functions.

PyTorch Testing (Python)
Python PyTorch
def train(args, model, device, train_loader, optimizer, epoch):
  for batch_idx, (data, target) in enumerate(train_loader):
    data, target =,
    output = model(data)
    loss = F.nll_loss(output, target)
    if batch_idx % args.log_interval == 0:
      print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
        epoch, batch_idx * len(data), len(train_loader.dataset),
        100. * batch_idx / len(train_loader), loss.item()))

def test(args, model, device, test_loader):
  test_loss = 0
  correct = 0
  with torch.no_grad():
    for data, target in test_loader:
      data, target =,
      output = model(data)
      # sum up batch loss
      test_loss += F.nll_loss(output, target, reduction='sum').item()
      # get the index of the max log-probability
      pred = output.argmax(dim=1, keepdim=True)
      correct += pred.eq(target.view_as(pred)).sum().item()

  test_loss /= len(test_loader.dataset)

  print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
    test_loss, correct, len(test_loader.dataset),
    100. * correct / len(test_loader.dataset)))

Set options and train.

PyTorch Testing (Python)
Python PyTorch
args = type('', (), {})()
args.batch_size = 64
args.test_batch_size = 1000
args.epochs = 10 = 0.01
args.momentum = 0.5
args.no_cuda = False
args.seed = 1
args.log_interval = 10
args.save_model = True

use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")

kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
train_loader =
  datasets.MNIST('../data', train=True, download=True,
                 transforms.Normalize((0.1307,), (0.3081,))
  batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader =
  datasets.MNIST('../data', train=False, transform=transforms.Compose([
                 transforms.Normalize((0.1307,), (0.3081,))
  batch_size=args.test_batch_size, shuffle=True, **kwargs)

model = Net().to(device)
optimizer = optim.SGD(model.parameters(),, momentum=args.momentum)

for epoch in range(1, args.epochs + 1):
  train(args, model, device, train_loader, optimizer, epoch)
  test(args, model, device, test_loader)

if (args.save_model):,"/results/")

Using a Pre-trained Model

Adapted from the fast-neural-style example.

Download the PyTorch example repo to get the relevant module files for the 'fast-neural-style' example. PYTHONPATH is set to the example's directory in the Runtime Settings, so import will work.

The pre-trained models are downloaded, saved to results, and locked. Unzip them.

wget --progress=dot:giga -O /results/ \
-d /


import argparse
import os
import sys
import time
import re

import numpy as np
import torch
from torch.optim import Adam
from import DataLoader
from torchvision import datasets
from torchvision import transforms
import torch.onnx

import utils
from transformer_net import TransformerNet
from vgg import Vgg16

The main styling function.

PyTorch Testing (Python)
Python PyTorch
def stylize(args):
  device = torch.device("cuda" if args.cuda else "cpu")

  content_image = utils.load_image(args.content_image, scale=args.content_scale)
  content_transform = transforms.Compose([
    transforms.Lambda(lambda x: x.mul(255))
  content_image = content_transform(content_image)
  content_image = content_image.unsqueeze(0).to(device)

  if args.model.endswith(".onnx"):
    output = stylize_onnx_caffe2(content_image, args)
    with torch.no_grad():
      style_model = TransformerNet()
      state_dict = torch.load(args.model)
      # remove saved deprecated running_* keys in InstanceNorm from the checkpoint
      for k in list(state_dict.keys()):
        if'in\d+\.running_(mean|var)$', k):
          del state_dict[k]
      if args.export_onnx:
        assert args.export_onnx.endswith(".onnx"), 
          "Export model file should end with .onnx"
        output = torch.onnx._export(style_model, content_image, 
        output = style_model(content_image).cpu()
  utils.save_image(args.output_image, output[0])

Helper function, used with ONNX model files.

def stylize_onnx_caffe2(content_image, args):
  Read ONNX model and run it using Caffe2

  assert not args.export_onnx

  import onnx
  import onnx_caffe2.backend

  model = onnx.load(args.model)

  prepared_backend = onnx_caffe2.backend.prepare(model, device='CUDA' if args.cuda else 'CPU')
  inp = {model.graph.input[0].name: content_image.numpy()}
  c2_out =[0]

  return torch.from_numpy(c2_out)

Let's make a Euphonium look fancy.

It's fast, so we might as well run through all four models.

args = type('', (), {})()
args.content_scale = None
args.cuda = 1
args.export_onnx = ""
args.content_image = 
for style in ["candy","udnie","rain_princess","mosaic"]: args.output_image = "/results/styled-{}.jpg".format(style) args.model = "/saved_models/{}.pth".format(style) stylize(args)


Build the PyTorch Environment

Install the dependencies via conda. Nvidia drivers and libraries are loaded by setting the NEXTJOURNAL_MOUNT_CUDA environment variable.

conda install -c defaults -c intel \
  mkl-include mkl-dnn \
  pyyaml typing
conda install -c pytorch cuda92 magma-cuda92

conda clean -qtipy


Install torch and torchvision via pip. This ensures conda doesn't force redundant installs of cudatoolkit and cudnn.

pip install -f \
  "torch==${TORCH_VERSION}+${TORCH_CUDA}" \


import torch