Training a PyTorch Model
Data
In this example, we will have two datapoints.
{% x_1 = [1,1], x_2 = [2,1] %}
{% y_1 = f(\vec{x}_1) = 1 %}
{% y_2 = f(\vec{x}_2) = 3 %}
In order to run the model, we then stack the training set into two matrices.
X = torch.tensor([[1.0,1.0],[2.0,1.0]])
y = torch.tensor([[1],[3]])
Running the Optimizer
network = nn.Sequential(
nn.Linear(2, 5),
nn.Tanh(),
nn.Linear(5,1)
)
opt = op.SGD(network.parameters(), lr=0.001)
mse = nn.MSELoss()
def train(model, opt, loss, X,y, epochs=1):
for epoch in range(epochs):
y2 = model(X)
loss1 = loss(y2, y)
loss1.backward()
opt.step()
opt.zero_grad()
pass
pass
train(network, opt, mse, X, y, 10000)
Full Sample Code
import torch
import torch.optim as op
import torch.nn as nn
network = nn.Sequential(
nn.Linear(2, 5),
nn.Tanh(),
nn.Linear(5,1)
)
X = torch.tensor([[1.0,1.0],[2.0,1.0]])
y = torch.tensor([[1],[3]])
test = network(X)
print(test)
#args = torch.tensor([1.0,1.0], requires_grad=True)
opt = op.SGD(network.parameters(), lr=0.001)
mse = nn.MSELoss()
def train(model, opt, loss, X,y, epochs=1):
for epoch in range(epochs):
y2 = model(X)
loss1 = loss(y2, y)
loss1.backward()
opt.step()
opt.zero_grad()
pass
pass
train(network, opt, mse, X, y, 10000)
test2 = network(X)
print(test2)
loss_tensor = mse(network(X),y)
#loss_value is the total loss
loss_value = loss_tensor.item()