File size: 956 Bytes
dbadc63 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | import torch
import torch.nn as nn
from dataset import get_data
from model import LinearModel
from matplotlib import pyplot as plt
def train():
x_train, y_train = get_data()
model = LinearModel()
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(),lr=0.01)
epochs = 1000
y_axis = []
for epoch in range(epochs):
pred = model(x_train)
loss = loss_fn(pred,y_train)
optimizer.zero_grad()
loss.backward()
optimizer.step()
y_axis.append(loss.item())
if epoch % 100 == 0:
print(f'Epoch {epoch}, Loss: {loss.item()}')
x_axis = range(epochs)
plt.plot(x_axis, y_axis)
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Training Loss over Epochs')
torch.save(model.state_dict(),'models/model.pt')
print('Model saved!')
plt.show()
if __name__ == '__main__':
train() |