-
Notifications
You must be signed in to change notification settings - Fork 0
/
train_autoencoder.py
116 lines (94 loc) · 3.33 KB
/
train_autoencoder.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import torch
from torch.utils.data import DataLoader
import numpy as np
from model.autoencoder import Autoencoder
from chess_dataset import ChessDataset
def train(model, train_data, val_data, epochs, loss_f, patient, lr, decay, save_path):
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
history = {"val_loss": [], "train_loss": []}
curr_patient = patient
# Train loop
for epoch in range(epochs):
running_loss = 0
model.train()
for batch_idx, states in enumerate(train_data):
optimizer.zero_grad()
reconstructed, _ = model(states)
loss = loss_f(reconstructed, states)
loss.backward()
optimizer.step()
# Calculate training statistics
running_loss += loss.item()
# Print every 1000
if batch_idx % 100 == 0:
print("Training epoch {}: {}/{}\t{:.0%}\tLoss: {:.3f}".format(epoch,
batch_idx,
len(train_data),
batch_idx/len(train_data),
loss.item()/len(states)))
val_loss = validation(model, val_data, loss_f)
# Decay
for p in optimizer.param_groups:
p["lr"] *= decay
# Save training data
history["train_loss"].append(running_loss/len(train_data.dataset))
history["val_loss"].append(val_loss)
# Save model every epochs
save(model, save_path, lr, decay, epoch)
# Early stopping with loss and patient epoch on validation result
if len(history["val_loss"]) > 1:
if val_loss - history["val_loss"][-2] > 0:
curr_patient -= 1
if curr_patient < 0:
print("Early stopping")
return history
return history
def validation(model, val_data, loss_f):
running_loss = 0
model.eval()
for states in val_data:
reconstructed, _ = model(states)
loss = loss_f(reconstructed, states)
running_loss += loss.item()
print("====> Test loss: {:.3f}".format(running_loss/len(val_data.dataset)))
return running_loss/len(val_data.dataset)
def save(model, path, lr, decay, epoch):
p = "{}/lr_{}_decay_{}.pt".format(path, lr, decay)
torch.save({
"epoch": epoch,
"model_state_dict": model.state_dict()
}, p)
print("Model saved on {}".format(p))
if __name__ == "__main__":
# Model
model = Autoencoder()
# Dataset
white_wins = np.load("./dataset/white/white_wins.npy", mmap_mode="c")
black_wins = np.load("./dataset/black/black_wins.npy", mmap_mode="c")
train_data = ChessDataset(white_wins, black_wins, train_split=0.8)
test_data = ChessDataset(white_wins, black_wins, train=False, train_split=0.8)
print(len(train_data), len(test_data))
trainloader = DataLoader(train_data, batch_size=128, shuffle=True)
testloader = DataLoader(test_data, batch_size=128, shuffle=False)
# Train
print("Start training")
epochs = 1
lr = 5e-4
decay = 0.95
save_path="./checkpoints/autoencoder"
loss_f = torch.nn.BCELoss(size_average=False)
resume = False
if resume:
c = torch.load("./checkpoints/autoencoder/lr_0.005_decay_0.95.pt")
model.load_state_dict(c["model_state_dict"])
his = train(model=model,
train_data=trainloader,
val_data=testloader,
epochs=epochs,
patient=2,
loss_f=loss_f,
lr=lr,
decay=decay,
save_path=save_path)
from utils import plot
plot(his)