86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
|
import numpy as np
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.optim as optim
|
|||
|
|
import torch.utils.data as data
|
|||
|
|
import torch
|
|||
|
|
from model import CharModel
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_data(filename="wonderland.txt"):
|
|||
|
|
# загружаем датасет и приводим к нижнему регистру
|
|||
|
|
filename = "wonderland.txt"
|
|||
|
|
raw_text = open(filename, 'r', encoding='utf-8').read()
|
|||
|
|
raw_text = raw_text.lower()
|
|||
|
|
|
|||
|
|
# делаем сопоставление текста с соответствующим ему значением
|
|||
|
|
chars = sorted(list(set(raw_text)))
|
|||
|
|
char_to_int = dict((c, i) for i, c in enumerate(chars))
|
|||
|
|
|
|||
|
|
# статистика обучаемых данных
|
|||
|
|
n_chars = len(raw_text)
|
|||
|
|
n_vocab = len(chars)
|
|||
|
|
print("Total Characters: ", n_chars)
|
|||
|
|
print("Total Vocab: ", n_vocab)
|
|||
|
|
|
|||
|
|
# подготовка датасета
|
|||
|
|
seq_length = 100
|
|||
|
|
dataX = []
|
|||
|
|
dataY = []
|
|||
|
|
for i in range(0, n_chars - seq_length, 1):
|
|||
|
|
seq_in = raw_text[i:i + seq_length]
|
|||
|
|
seq_out = raw_text[i + seq_length]
|
|||
|
|
dataX.append([char_to_int[char] for char in seq_in])
|
|||
|
|
dataY.append(char_to_int[seq_out])
|
|||
|
|
n_patterns = len(dataX)
|
|||
|
|
print("Total Patterns: ", n_patterns)
|
|||
|
|
|
|||
|
|
# --- переводим данные к тензору, чтобы рабоать с ними внутри pytorch ---
|
|||
|
|
X = torch.tensor(dataX, dtype=torch.float32).reshape(n_patterns, seq_length, 1)
|
|||
|
|
X = X / float(n_vocab)
|
|||
|
|
y = torch.tensor(dataY)
|
|||
|
|
print(X.shape, y.shape)
|
|||
|
|
|
|||
|
|
return X, y, char_to_int
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
X, y, char_to_int = get_data()
|
|||
|
|
|
|||
|
|
n_epochs = 40
|
|||
|
|
batch_size = 128
|
|||
|
|
model = CharModel()
|
|||
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|||
|
|
print(f"device: {device}")
|
|||
|
|
model.to(device)
|
|||
|
|
|
|||
|
|
optimizer = optim.Adam(model.parameters())
|
|||
|
|
loss_fn = nn.CrossEntropyLoss(reduction="sum")
|
|||
|
|
loader = data.DataLoader(data.TensorDataset(X, y), shuffle=True, batch_size=batch_size)
|
|||
|
|
|
|||
|
|
best_model = None
|
|||
|
|
best_loss = np.inf
|
|||
|
|
|
|||
|
|
for epoch in range(n_epochs):
|
|||
|
|
model.train()
|
|||
|
|
for X_batch, y_batch in loader:
|
|||
|
|
y_pred = model(X_batch.to(device))
|
|||
|
|
loss = loss_fn(y_pred, y_batch.to(device))
|
|||
|
|
|
|||
|
|
optimizer.zero_grad()
|
|||
|
|
loss.backward()
|
|||
|
|
optimizer.step()
|
|||
|
|
|
|||
|
|
# Validation
|
|||
|
|
model.eval()
|
|||
|
|
loss = 0
|
|||
|
|
with torch.no_grad():
|
|||
|
|
for X_batch, y_batch in loader:
|
|||
|
|
y_pred = model(X_batch.to(device))
|
|||
|
|
loss += loss_fn(y_pred, y_batch.to(device))
|
|||
|
|
if loss < best_loss:
|
|||
|
|
best_loss = loss
|
|||
|
|
best_model = model.state_dict()
|
|||
|
|
print("Epoch %d: Cross-entropy: %.4f" % (epoch, loss))
|
|||
|
|
|
|||
|
|
torch.save([best_model, char_to_int], "single-char.pth")
|
|||
|
|
|