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