pytorch1

A single pytorch model based on the FastionMNIST example given in the PyTorch documentation runs with input from a files and output to a file.

PyTorch Version

Model Code:

  1# Example from https://pytorch.org/tutorials/beginner/basics/
  2#   quickstart_tutorial.html
  3import os
  4from torch import nn
  5weights_file = os.path.normpath(
  6    os.path.join(
  7        os.path.dirname(__file__), 'model_weights.pth'))
  8data_root = os.path.normpath(
  9    os.path.join(os.path.dirname(__file__), '..', 'Input', 'data'))
 10input_file = os.path.normpath(
 11    os.path.join(os.path.dirname(__file__), '..', 'Input', 'input.png'))
 12output_file = os.path.normpath(
 13    os.path.join(os.path.dirname(__file__), '..', 'Output', 'output.txt'))
 14
 15
 16class NeuralNetwork(nn.Module):
 17    def __init__(self):
 18        super().__init__()
 19        self.flatten = nn.Flatten()
 20        self.linear_relu_stack = nn.Sequential(
 21            nn.Linear(28 * 28, 512),
 22            nn.ReLU(),
 23            nn.Linear(512, 512),
 24            nn.ReLU(),
 25            nn.Linear(512, 10)
 26        )
 27
 28    def forward(self, x):
 29        x = self.flatten(x)
 30        logits = self.linear_relu_stack(x)
 31        return logits
 32
 33
 34def input_transform(x):
 35    from torchvision.transforms import ToTensor
 36    from PIL import Image
 37    return (ToTensor()(Image.fromarray(x)), )
 38
 39
 40def output_transform(x):
 41    out = x.detach().numpy()
 42    out = out.astype({'names': ['f0'],
 43                      'formats': [out.dtype]})
 44    out = out.flatten()
 45    return out
 46
 47
 48def load_model():
 49    import torch
 50    import numpy as np
 51    from torchvision import datasets
 52    from torchvision.transforms import ToTensor, functional
 53    
 54    test_data = datasets.FashionMNIST(
 55        root=data_root,
 56        train=False,
 57        download=True,
 58        transform=ToTensor()
 59    )
 60    
 61    model = NeuralNetwork()
 62    model.load_state_dict(torch.load(weights_file))
 63    model.eval()
 64
 65    mock = test_data[0][0]
 66    im = functional.to_pil_image(mock)
 67    im.save(input_file)
 68    out = model(mock)
 69    arr = out.detach().numpy()
 70    np.savetxt(output_file, arr)
 71
 72
 73def train_model():
 74    import torch
 75    from torch.utils.data import DataLoader
 76    from torchvision import datasets
 77    from torchvision.transforms import ToTensor
 78    from torch import nn
 79    training_data = datasets.FashionMNIST(
 80        root=data_root,
 81        train=True,
 82        download=True,
 83        transform=ToTensor()
 84    )
 85    test_data = datasets.FashionMNIST(
 86        root=data_root,
 87        train=False,
 88        download=True,
 89        transform=ToTensor()
 90    )
 91    train_dataloader = DataLoader(
 92        training_data, batch_size=64, shuffle=True)
 93    test_dataloader = DataLoader(
 94        test_data, batch_size=64, shuffle=True)
 95
 96    model = NeuralNetwork()
 97
 98    loss_fn = nn.CrossEntropyLoss()
 99    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
100
101    def train(dataloader, model, loss_fn, optimizer):
102        size = len(dataloader.dataset)
103        model.train()
104        for batch, (X, y) in enumerate(dataloader):
105
106            # Compute prediction error
107            pred = model(X)
108            loss = loss_fn(pred, y)
109
110            # Backpropagation
111            loss.backward()
112            optimizer.step()
113            optimizer.zero_grad()
114
115            if batch % 100 == 0:
116                loss, current = loss.item(), (batch + 1) * len(X)
117                print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
118
119    def test(dataloader, model, loss_fn):
120        size = len(dataloader.dataset)
121        num_batches = len(dataloader)
122        model.eval()
123        test_loss, correct = 0, 0
124        with torch.no_grad():
125            for X, y in dataloader:
126                pred = model(X)
127                test_loss += loss_fn(pred, y).item()
128                correct += (
129                    pred.argmax(1) == y).type(torch.float).sum().item()
130        test_loss /= num_batches
131        correct /= size
132        print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, "
133              f"Avg loss: {test_loss:>8f} \n")
134
135    epochs = 5
136    for t in range(epochs):
137        print(f"Epoch {t+1}\n-------------------------------")
138        train(train_dataloader, model, loss_fn, optimizer)
139        test(test_dataloader, model, loss_fn)
140        print("Done!")
141
142    torch.save(model.state_dict(), weights_file)
143    test(test_dataloader, model, loss_fn)

Model YAML:

 1model:
 2  name: model
 3  language: pytorch
 4  args: ./src/model.py:NeuralNetwork
 5  weights: ./src/model_weights.pth
 6  input_transform: ./src/model.py:input_transform
 7  output_transform: ./src/model.py:output_transform
 8  inputs:
 9    - name: input
10      default_file:
11        name: ./Input/input.png
12        filetype: png
13  outputs:
14    - name: output
15      as_array: true
16      default_file:
17        name: ./output.txt
18        filetype: table