zachzhang / deepspeech2 Goto Github PK

model.py
def forward(self,x):
h0 = Variable(torch.ones(2, x.size()[0], self.h))
if torch.cuda.is_available():
h0 = h0.cuda()
x = seqWise(self.bn , x)
x = self.rnn(x,h0)[0]
x = x[:, :, :self.h] + x[:, :,self.h:]

train.py
if cuda:
model = model.cuda()
print('model.cuda() is over')

criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),lr = 1)
for epoch in range(epochs):
for i in range(num_batches):
X, y, y_sizes = get_batch(df, batch_size=batch_size)
y = [eval(_) for _ in y] #old y:list of str
# print(len(X),len(y),len(y_sizes),y_sizes,(len(y[0]),len(y[1])),type(y[0]),sep ='\n')
X = Variable(torch.from_numpy(np.expand_dims(X,1)), requires_grad=False ).float()

    #自定义 y
    y = [sum(_)/len(_) for _ in y]
    y = np.array(y).mean()
    y = np.array([y])
    y = Variable(torch.from_numpy(y), requires_grad=False).float()/100
    if cuda:
        X = X.cuda()
        y = y.cuda()
    out = model(X)
    out = out.mean()
    print('out',out,'target',y,sep=' ')
    loss = criterion(out, y)
    
    optimizer.zero_grad()
    loss.backward()
    # SGD step
    optimizer.step()

    if cuda:
        torch.cuda.synchronize()
print('--------epoch',epoch,'loss',loss,sep=' ')

print('train is over final')