103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
import torch.optim as optim
|
||
|
|
from torchvision import transforms
|
||
|
|
import pickle
|
||
|
|
import numpy as np
|
||
|
|
import pandas as pd
|
||
|
|
from word2vec import Word2Vec
|
||
|
|
|
||
|
|
class FFN(nn.Module):
|
||
|
|
|
||
|
|
def __init__(self, input_dim, output_dim, hidden1_size, hidden2_size, lr, epochs, batch_size):
|
||
|
|
super(FFN, self).__init__()
|
||
|
|
self.path = 'model1.pickle'
|
||
|
|
self.lr = lr
|
||
|
|
self.epochs = epochs
|
||
|
|
self.output_dim = output_dim
|
||
|
|
self.word2vec = Word2Vec()
|
||
|
|
self.word2vec.load()
|
||
|
|
self.batch_size = batch_size
|
||
|
|
self.input_dim = input_dim
|
||
|
|
self.fc1 = nn.Linear(batch_size, hidden1_size)
|
||
|
|
self.fc2 = nn.Linear(hidden1_size, hidden2_size)
|
||
|
|
self.fc3 = nn.Linear(hidden2_size, hidden2_size)
|
||
|
|
self.fc4 = nn.Linear(hidden2_size, hidden2_size)
|
||
|
|
self.fc5 = nn.Linear(hidden2_size, batch_size)
|
||
|
|
|
||
|
|
def forward(self, data):
|
||
|
|
data = F.relu(self.fc1(data))
|
||
|
|
data = F.relu(self.fc2(data))
|
||
|
|
data = F.relu(self.fc3(data))
|
||
|
|
data = F.relu(self.fc4(data))
|
||
|
|
data = F.sigmoid(self.fc5(data))
|
||
|
|
return data
|
||
|
|
|
||
|
|
def serialize(self):
|
||
|
|
with open(self.path, 'wb') as file:
|
||
|
|
pickle.dump(self, file)
|
||
|
|
|
||
|
|
def load(self):
|
||
|
|
with open(self.path, 'rb') as file:
|
||
|
|
self = pickle.load(file)
|
||
|
|
|
||
|
|
def batch(self, iterable, n=1):
|
||
|
|
l = len(iterable)
|
||
|
|
for ndx in range(0, l, n):
|
||
|
|
yield iterable[ndx:min(ndx + n, l)]
|
||
|
|
|
||
|
|
"""
|
||
|
|
data is a tuple of embedding vector and a label of 0/1
|
||
|
|
"""
|
||
|
|
def train(self, data, expected):
|
||
|
|
self.zero_grad()
|
||
|
|
criterion = torch.nn.BCELoss()
|
||
|
|
optimizer = optim.Adam(self.parameters(), lr=self.lr)
|
||
|
|
batch_size = self.batch_size
|
||
|
|
num_of_classes = self.output_dim
|
||
|
|
for epoch in range(self.epochs):
|
||
|
|
epoch_loss = 0.0
|
||
|
|
idx = 0
|
||
|
|
for i in range(0, int(len(data)/batch_size)*batch_size, batch_size):
|
||
|
|
inputs = data[i:i + batch_size]
|
||
|
|
labels = expected[i:i+ batch_size]
|
||
|
|
optimizer.zero_grad()
|
||
|
|
outputs = self.forward(torch.tensor(self.word2vec.list_of_sentences2vec(inputs)))
|
||
|
|
target = torch.tensor(labels.values).double()
|
||
|
|
loss = criterion(outputs.view(batch_size), target.view(-1,))
|
||
|
|
loss.backward()
|
||
|
|
optimizer.step()
|
||
|
|
|
||
|
|
epoch_loss += loss.item()
|
||
|
|
if(idx % 1000 == 0):
|
||
|
|
print('epoch: {}, idx: {}, loss: {}'.format(epoch, idx, epoch_loss/1000))
|
||
|
|
epoch_loss = 0
|
||
|
|
idx += 1
|
||
|
|
self.serialize()
|
||
|
|
|
||
|
|
def test(self, data, expected, path):
|
||
|
|
correct = 0
|
||
|
|
incorrect = 0
|
||
|
|
total = 0
|
||
|
|
predictions = []
|
||
|
|
batch_size = self.batch_size
|
||
|
|
for i in range(0, int(len(data)/batch_size)*batch_size, batch_size):
|
||
|
|
inputs = data[i:i + batch_size]
|
||
|
|
labels = expected[i:i+ batch_size]
|
||
|
|
predicted = self.forward(torch.tensor(self.word2vec.list_of_sentences2vec(inputs)))
|
||
|
|
score = [1 if x > 0.5 else 0 for x in predicted]
|
||
|
|
|
||
|
|
for x, y in zip(score, labels):
|
||
|
|
if(x == y):
|
||
|
|
correct += 1
|
||
|
|
else:
|
||
|
|
incorrect += 1
|
||
|
|
predictions.append(score)
|
||
|
|
|
||
|
|
print(correct)
|
||
|
|
print(incorrect)
|
||
|
|
print(correct/(incorrect + correct))
|
||
|
|
df = pd.DataFrame(np.asarray(predictions).reshape(int(len(data)/batch_size)*batch_size))
|
||
|
|
df.reset_index(drop=True, inplace=True)
|
||
|
|
df.to_csv(path, sep="\t", index=False)
|