-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
40 lines (33 loc) · 1.24 KB
/
train.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
import hiddenlayer as hl
import torch
from torch import nn
import torch.optim as optim
from model_design import AutoEncoder
from data import get_dataset
if __name__ == '__main__':
train_loader, train_data_y, test_loader, test_data_y = get_dataset(train_batch_size=128)
lr = 3e-3
auto_encoder_model = AutoEncoder().cuda()
optimizer = optim.Adam(auto_encoder_model.parameters(), lr=lr)
criterion = nn.MSELoss()
history = hl.History()
canvas = hl.Canvas()
print("Training...")
for epoch in range(10):
train_loss_epoch = 0
train_num = 0
for step, images in enumerate(train_loader):
images = images.cuda()
_, output = auto_encoder_model(images)
loss = criterion(output, images)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss_epoch += loss.item()
train_num = train_num + images.size(0)
train_loss = train_loss_epoch / train_num
history.log(epoch, train_loss=train_loss)
with canvas:
canvas.draw_plot(history['train_loss'])
torch.save({"state_dict": auto_encoder_model.state_dict(),}, "./checkpoint/auto_encoder.pth")
print('Save model done!')