BaseTen
Search…
PyTorch
PyTorch models are handled in BaseTen just like scikit-learn and TensorFlow models, but require additional file(s) defining the model class.
If your model class MyModel is defined in the file my_model.py, add the following keyword argument to the baseten.deploy call:
main.py
my_model.py
1
import baseten
2
baseten.deploy(
3
my_model, # e.g: a PyTorch model MyModel
4
model_name='My pytorch model',
5
model_files=['my_model.py'],
6
)
Copied!
1
import torch
2
import torch.nn as nn
3
import torch.nn.functional as F
4
5
class MyModel(nn.Module):
6
def __init__(self):
7
super(Net, self).__init__()
8
self.conv1 = nn.Conv2d(1, 32, 3, 1)
9
self.conv2 = nn.Conv2d(32, 64, 3, 1)
10
self.dropout1 = nn.Dropout(0.25)
11
self.dropout2 = nn.Dropout(0.5)
12
self.fc1 = nn.Linear(9216, 128)
13
self.fc2 = nn.Linear(128, 10)
14
15
def forward(self, x):
16
x = self.conv1(x)
17
x = F.relu(x)
18
x = self.conv2(x)
19
x = F.relu(x)
20
x = F.max_pool2d(x, 2)
21
x = self.dropout1(x)
22
x = torch.flatten(x, 1)
23
x = self.fc1(x)
24
x = F.relu(x)
25
x = self.dropout2(x)
26
x = self.fc2(x)
27
output = F.log_softmax(x, dim=1)
28
return output
Copied!
baseten.deploy will deploy your model to BaseTen and print out a URL. Go there on your browser to see its deployment status and other useful information.
Last modified 6mo ago
Copy link