-
Notifications
You must be signed in to change notification settings - Fork 5
/
train_net.py
55 lines (40 loc) · 1.56 KB
/
train_net.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
import os
import torch.nn as nn
import torch.nn.init as init
import torch.optim as optim
import torch
import torchvision
from define_net import Net
from torch.autograd import Variable
from paragraph_set import Paragraph_set
from dataloader_modified import DataLoader
if __name__ == '__main__':
path_ = os.path.abspath('.')
trainset = Paragraph_set(path_+'/hongloumeng/')
print trainset.get_paragraph_size()
print trainset.get_word_size()
# the length of samples here are different, so we can't directly use DataLoader provided by PyTorch
trainloader = DataLoader(trainset,batch_size=8,shuffle=True,num_workers=2)
net = Net(trainset.get_word_size())
print net
criterion = nn.NLLLoss()
optimizer = optim.RMSprop(net.parameters(), lr=0.01, weight_decay=0.0001)
for epoch in range(5): #
running_loss = 0.0
for i,batch in enumerate(trainloader,0):
loss = 0
optimizer.zero_grad()
for data in batch:
inputs,targets = data
inputs,targets = Variable(inputs),Variable(targets)
hidden = net.initHidden()
outputs,hidden = net(inputs,hidden)
loss = loss + criterion(outputs,targets)
loss.backward()
optimizer.step()
running_loss += loss.data[0]
if i%10 == 9:
print('[%d, %4d] loss: %.3f' % (epoch+1,i+1,running_loss/80)) # step is 10 and batch_size is 8
running_loss = 0.0
print('Finished Training')
torch.save(net.state_dict(),path_+'/net.pth')