68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
import torch, numpy as np
|
|
|
|
class NeuralNetworkModel(torch.nn.Module):
|
|
def __init__(self, features=100):
|
|
super(NeuralNetworkModel, self).__init__()
|
|
self.fc1 = torch.nn.Linear(features, 500)
|
|
self.fc2 = torch.nn.Linear(500, 1)
|
|
|
|
def forward(self, x):
|
|
x = self.fc1(x)
|
|
x = torch.relu(x)
|
|
x = self.fc2(x)
|
|
x = torch.sigmoid(x)
|
|
return x
|
|
|
|
def getMetrics(model, X_dataset, Y_dataset, criterion, batchSize):
|
|
loss_score = 0
|
|
acc_score = 0
|
|
items_total = 0
|
|
model.eval()
|
|
for i in range(0, Y_dataset.shape[0], batchSize):
|
|
X = X_dataset[i:i+batchSize]
|
|
X = torch.tensor(X.astype(np.float32))
|
|
Y = Y_dataset[i:i+batchSize]
|
|
Y = torch.tensor(Y.astype(np.float32)).reshape(-1,1)
|
|
Y_predictions = model(X)
|
|
acc_score += torch.sum((Y_predictions > 0.5) == Y).item()
|
|
items_total += Y.shape[0]
|
|
|
|
loss = criterion(Y_predictions, Y)
|
|
|
|
loss_score += loss.item() * Y.shape[0]
|
|
return (loss_score / items_total), (acc_score / items_total)
|
|
|
|
def trainModel(model, trainX, trainY, devX, devY, optimizer, criterion=torch.nn.BCELoss(), epochs=5, batchSize=256):
|
|
for epoch in range(epochs):
|
|
loss_score = 0
|
|
acc_score = 0
|
|
items_total = 0
|
|
model.train()
|
|
for i in range(0, trainY.shape[0], batchSize):
|
|
X = trainX[i:i+batchSize]
|
|
X = torch.tensor(X.astype(np.float32))
|
|
Y = trainY[i:i+batchSize]
|
|
Y = torch.tensor(Y.astype(np.float32)).reshape(-1,1)
|
|
Y_predictions = model(X)
|
|
acc_score += torch.sum((Y_predictions > 0.5) == Y).item()
|
|
items_total += Y.shape[0]
|
|
|
|
optimizer.zero_grad()
|
|
loss = criterion(Y_predictions, Y)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
loss_score += loss.item() * Y.shape[0]
|
|
|
|
print(f'Epoch {epoch+1}/{epochs}')
|
|
loss, accuracy = getMetrics(model, trainX, trainY, criterion, batchSize)
|
|
print(f'Train set\nloss = {loss}, accuracy = {accuracy}')
|
|
|
|
def flatten(t):
|
|
return [str(int(item)) for sublist in t for item in sublist]
|
|
|
|
def predict(model, testX):
|
|
testX = torch.tensor(testX.astype(np.float32))
|
|
with torch.no_grad():
|
|
return flatten(model(testX).round().tolist())
|