49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
class Model(nn.Module):
|
||
|
|
"""NN Model class"""
|
||
|
|
|
||
|
|
def __init__(self, input_dim=300, hidden_dim=600, output_dim=1):
|
||
|
|
"""Initializes new instance of Model class"""
|
||
|
|
|
||
|
|
super(Model, self).__init__()
|
||
|
|
|
||
|
|
self.input_dim = input_dim
|
||
|
|
self.hidden_dim = hidden_dim
|
||
|
|
self.output_dim = output_dim
|
||
|
|
|
||
|
|
self.fc1 = nn.Linear(self.input_dim, self.hidden_dim)
|
||
|
|
self.fc2 = nn.Linear(self.hidden_dim, self.output_dim)
|
||
|
|
|
||
|
|
self.criterion = nn.BCELoss()
|
||
|
|
self.optimizer = torch.optim.SGD(self.parameters(), lr=0.01)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
"""Step forward learning fn"""
|
||
|
|
|
||
|
|
x = self.fc1(x)
|
||
|
|
x = torch.relu(x)
|
||
|
|
x = self.fc2(x)
|
||
|
|
x = torch.sigmoid(x)
|
||
|
|
return x
|
||
|
|
|
||
|
|
def run_training(self, X_train, Y_train, batch_size, epochs_count):
|
||
|
|
for _ in range(epochs_count):
|
||
|
|
self.train()
|
||
|
|
for i in range(0, Y_train.shape[0], batch_size):
|
||
|
|
X = X_train[i: i + batch_size]
|
||
|
|
X = torch.tensor(X)
|
||
|
|
y = Y_train[i: i + batch_size]
|
||
|
|
y = torch.tensor(
|
||
|
|
y.astype(np.float32).to_numpy()).reshape(-1, 1)
|
||
|
|
|
||
|
|
outputs = self(X.float())
|
||
|
|
loss = self.criterion(outputs, y)
|
||
|
|
|
||
|
|
self.optimizer.zero_grad()
|
||
|
|
loss.backward()
|
||
|
|
self.optimizer.step()
|