77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
import torch
|
|
from torch import nn
|
|
from torch import nn, optim
|
|
from torch.utils.data import DataLoader
|
|
import numpy as np
|
|
import random
|
|
|
|
import shark
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
self.lstm = nn.LSTM(
|
|
input_size=8,
|
|
hidden_size=16,
|
|
num_layers=3,
|
|
dropout=0.1,
|
|
)
|
|
self.fc = nn.Linear(16, 1)
|
|
self.out = nn.Sigmoid()
|
|
|
|
def forward(self, x, prev_state):
|
|
output, state = self.lstm(x, prev_state)
|
|
logits = self.fc(output)
|
|
val = self.out(logits)
|
|
#print(str(logits.item())+" > "+str(val.item()))
|
|
return val, state
|
|
|
|
def init_state(self, sequence_length):
|
|
return (torch.zeros(3, 1, 16),
|
|
torch.zeros(3, 1, 16))
|
|
|
|
def train(model, seq_len=16*64):
|
|
tid = str(int(random.random()*99999)).zfill(5)
|
|
print("[i] I am "+str(tid))
|
|
ltLoss = 50
|
|
lltLoss = 51
|
|
model.train()
|
|
|
|
criterion = nn.BCELoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=0.0001)
|
|
|
|
state_h = [None,None]
|
|
state_c = [None,None]
|
|
blob = [None,None]
|
|
correct = [None,None]
|
|
|
|
for epoch in range(1024):
|
|
state_h[0], state_c[0] = model.init_state(seq_len)
|
|
state_h[1], state_c[1] = model.init_state(seq_len)
|
|
|
|
blob[0], _ = shark.getSample(min(seq_len, 16*(epoch+1)), 0)
|
|
blob[1], _ = shark.getSample(min(seq_len, 16*(epoch+1)), 1)
|
|
optimizer.zero_grad()
|
|
for i in range(len(blob[0])):
|
|
for t in range(2):
|
|
x = torch.tensor([[[float(d) for d in bin(blob[t][i])[2:].zfill(8)]]], dtype=torch.float32)
|
|
y_pred, (state_h[t], state_c[t]) = model(x, (state_h[t], state_c[t]))
|
|
loss = criterion(y_pred[0][0][0], torch.tensor(t, dtype=torch.float32))
|
|
|
|
state_h[t] = state_h[t].detach()
|
|
state_c[t] = state_c[t].detach()
|
|
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
correct[t] = round(y_pred.item()) == t
|
|
ltLoss = ltLoss*0.9 + 0.1*loss.item()
|
|
lltLoss = lltLoss*0.9 + 0.1*ltLoss
|
|
print({ 'epoch': epoch, 'loss': loss.item(), 'ltLoss': ltLoss, 'ok0': correct[0], 'ok1': correct[1], 'succ': correct[0] and correct[1] })
|
|
if epoch % 8 == 0:
|
|
torch.save(model.state_dict(), 'model_savepoints/'+tid+'_'+str(epoch)+'.n')
|
|
|
|
model = Model()
|
|
|
|
train(model)
|