PyTorch Environment

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

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)

def train(args, model, device, train_loader, optimizer, epoch):

  model.train()

  for batch_idx, (data, target) in enumerate(train_loader):

    data, target = data.to(device), target.to(device)

    optimizer.zero_grad()

    output = model(data)

    loss = F.nll_loss(output, target)

    loss.backward()

    optimizer.step()

    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):

  model.eval()

  test_loss = 0

  correct = 0

  with torch.no_grad():

    for data, target in test_loader:

      data, target = data.to(device), target.to(device)

      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)))

args = type('', (), {})()

args.batch_size = 64

args.test_batch_size = 1000

args.epochs = 10

args.lr = 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()

torch.manual_seed(args.seed)

device = torch.device("cuda" if use_cuda else "cpu")

kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

train_loader = torch.utils.data.DataLoader(

  datasets.MNIST('../data', train=True, download=True,

                 transform=transforms.Compose([

                 transforms.ToTensor(),

                 transforms.Normalize((0.1307,), (0.3081,))

                 ])),

  batch_size=args.batch_size, shuffle=True, **kwargs)

test_loader = torch.utils.data.DataLoader(

  datasets.MNIST('../data', train=False, transform=transforms.Compose([

                 transforms.ToTensor(),

                 transforms.Normalize((0.1307,), (0.3081,))

                 ])),

  batch_size=args.test_batch_size, shuffle=True, **kwargs)

model = Net().to(device)

optimizer = optim.SGD(model.parameters(), lr=args.lr, 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):

  torch.save(model.state_dict(),"/results/mnist_cnn.pt")

wget --progress=dot:giga -O /results/saved_models.zip \

  https://www.dropbox.com/s/lrvwfehqdcxoza8/saved_models.zip?dl=1

unzip 
saved_models.zip -d /

import argparse

import os

import sys

import time

import re

import numpy as np

import torch

from torch.optim import Adam

from torch.utils.data import DataLoader

from torchvision import datasets

from torchvision import transforms

import torch.onnx

import utils

from transformer_net import TransformerNet

from vgg import Vgg16

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.ToTensor(),

    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)

  else:

    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 re.search(r'in\d+\.running_(mean|var)$', k):

          del state_dict[k]

      style_model.load_state_dict(state_dict)

      style_model.to(device)

      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, 

                                    args.export_onnx).cpu()

      else:

        output = style_model(content_image).cpu()

  utils.save_image(args.output_image, output[0])

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 = prepared_backend.run(inp)[0]

  return torch.from_numpy(c2_out)

args = type('', (), {})()

args.content_scale = None

args.cuda = 1

args.export_onnx = ""

args.content_image = 
Euphonium_Boosey_and_hawkes.jpg

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)

conda install -c defaults -c intel \

  mkl-include mkl-dnn \

  pyyaml typing

conda install -c pytorch cuda100 magma-cuda102

conda clean -qtipy

ldconfig

pip install \

  -f https://download.pytorch.org/whl/cu102/torch_stable.html \

  torch torchvision torchcsprng torchaudio torchtext \

  torch_model_archiver torchserve

import torch

print(torch.cuda.current_device())

print(torch.cuda.device(0))

print(torch.cuda.device_count())

print(torch.cuda.get_device_name(0))

print(torch.cuda.is_available())

print(torch.rand(100,100).cuda())

import torchvision.ops

conda list torch

pip list

nvidia-smi