lstmRearch/copyTask.py at master · backspacetg/lstmRearch · GitHub

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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
import lstmTest
import numpy as np
import torch as t
from torch import optim
import torch.nn as nn
import math
import visdom


t.manual_seed(1000)
T = 20
trainSize = 100000
testSize = 100
inputSize = 1
outputSize = 9
hideSize = 32
batchSize = 100
def dataGenerator(T):
    item = [1,2,3,4,5,6,7,8,0,9]
    ranks = np.random.randint(8,size=10)
    inputs = []
    #必须是一维tensor的数组，否则没有办法使用torch.cat函数
    for r in ranks:
        inputs.append([item[r]])
    for i in range(T-1):
        inputs.append([item[8]])
    inputs.append([item[9]])
    for j in range(10):
        inputs.append([item[8]])
    outputs = []
    for ii in range(T+10):
        outputs.append([item[8]])
    for rr in ranks:
        outputs.append([item[rr]])

    return t.FloatTensor(inputs), t.LongTensor(outputs) #返回longTensor原因是交叉熵损失函数的标签项只能是long的

def datasetGenerator(T,size):
    data = []
    labels = []
    for i in range(size):
        x,y = dataGenerator(T)
        data.append(x)
        labels.append(y)
    data = t.stack(data)
    labels = t.stack(labels)
    return data,labels.long()


class Net(nn.Module):
    def __init__(self,inputSize,hideSize,outputSize):
        super(Net,self).__init__()
        self.lstm = lstmTest.LstmCell(inputSize,hideSize)
        self.fc = nn.Linear(hideSize,outputSize)
    def forward(self,x,h,c):
        hideState,cellState = self.lstm(x,h,c)
        output = self.fc(hideState)
        return output,hideState,cellState


if __name__ == '__main__':
    t.set_num_threads(8)
    device = t.device("cuda:0" if t.cuda.is_available() else "cpu")
    net = Net(inputSize,hideSize,outputSize).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(),lr=0.001)
    #加载训练结果
    #net.load_state_dict(t.load('save/net2.pkl'))
    #ptimizer.load_state_dict(t.load('save/optimizer2.pkl'))
    #生成训练集
    trainInput, supposedOutputs = datasetGenerator(T,trainSize)
    trainInput = trainInput.to(device)
    supposedOutputs = supposedOutputs.to(device)
    #生成测试集
    testInputs,testOutputs = datasetGenerator(T,testSize)
    testInputs = testInputs.to(device)
    testOutputs = testOutputs.to(device)
    #可视化
    #vis = visdom.Visdom(env=u'copyTask') #指定字符串为Unicode对象
    #vis.close()

    for epoch in range(90): #遍历数据集
        print('in epoch ',epoch)
        for i in range(trainSize//batchSize):
            loss = 0
            h = t.zeros(batchSize,hideSize)
            c = t.zeros(batchSize,hideSize)
            ind = np.random.choice(trainSize,batchSize)
            inputsquence = trainInput[ind]
            label = supposedOutputs[ind]
            inputsquence.transpose_(0,1)
            label.transpose_(0,1)
            for j in range(T+20):
                output,h,c = net(inputsquence[j],h,c)
                loss += criterion(output,label[j].squeeze(1))#由损失函数的输入格式决定

                if(math.isnan(loss)):
                    print(inputsquence.squeeze())
                    print(label.squeeze())
                    print(output,h,c)
                    print('lstm infomation: \n',net.lstm.combinedLinear.weight,'\n',net.lstm.combinedLinear.bias)
                    print('Linear infomation: \n',net.fc.weight,'\n',net.fc.bias)
                    quit()

            loss /= (T+20)
            net.zero_grad()
            loss.backward()
            optimizer.step()

            if(i%100==0):
                print('till train ',(trainSize//batchSize)*epoch+i,': loss=',loss)
                x= (trainSize//batchSize) * epoch + i
                #logx = math.log10(x+1)
                #vis.line(X=t.Tensor([logx]),Y=loss.unsqueeze(0),win='loss2',update='append')
                #vis.line(X=t.Tensor([x]), Y=loss.unsqueeze(0), win='loss', update='append')

    correctNumber = 0
    #正确率测试
    for jk in range(testSize):
        #print("result for test",jk)
        h = t.zeros(1,hideSize)
        c = t.zeros(1,hideSize)
        for testTime in range(T+20):
            output,h,c = net(testInputs[jk][testTime].unsqueeze(0),h,c)
            _,max = t.max(output,1)
            if t.equal(max.data,testOutputs[jk][testTime]):
                correctNumber+=1
    print('correctRate: ',correctNumber/(testSize*(T+20)))

    #保存模型
    t.save(net.state_dict(),'save/net0_95.pkl')
    t.save(optimizer.state_dict(),'save/optimizer0_95.pkl')