Tutorial 4: Exporting your modelΒΆ

After you have trained your model and you are satisfied with the weights, exporting your model back to the original library is straight-forward:

npu.export(model, PATH)

Where model is the output of npu.train() or the id of the model you want to export from your dashboard and PATH is the directory where your model will be exported to.

The following code is an example of exporting your torch model after being trained with the NPU API.

import os

import torch.nn as nn

import npu
import npu.vision.datasets as dset

npu.api(API_TOKEN)

PATH = os.getcwd() + '/'


class Net(nn.Module):
    def __init__(self):
        nn.Module.__init__(self)
        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2))
        self.fc1 = nn.Linear(in_features=32*6*6, out_features=10)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.flatten(start_dim=1)
        x = self.fc1(x)
        return x


model = Net()

model = npu.compile(model, input_shape=[1, 28, 28])
model_trained = npu.train(model,
                          train_data=dset.MNIST.train,
                          val_data=dset.MNIST.val,
                          loss=npu.loss.SparseCrossEntropyLoss,
                          optim=npu.optim.SGD(lr=0.001, momentum=0.001),
                          batch_size=100,
                          epochs=4,
                          asynchronous=True)

npu.export(model_trained, PATH)