|
47 | 47 | # convert test data into Variable, pick 2000 samples to speed up testing |
48 | 48 | test_data = dsets.MNIST(root='./mnist/', train=False, transform=transforms.ToTensor()) |
49 | 49 | test_x = test_data.test_data.type(torch.FloatTensor)[:2000]/255. # shape (2000, 28, 28) value in range(0,1) |
50 | | -test_y = test_data.test_labels.numpy().squeeze()[:2000] # covert to numpy array |
| 50 | +test_y = test_data.test_labels.numpy()[:2000] # covert to numpy array |
51 | 51 |
|
52 | 52 |
|
53 | 53 | class RNN(nn.Module): |
@@ -94,13 +94,13 @@ def forward(self, x): |
94 | 94 |
|
95 | 95 | if step % 50 == 0: |
96 | 96 | test_output = rnn(test_x) # (samples, time_step, input_size) |
97 | | - pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze() |
| 97 | + pred_y = torch.max(test_output, 1)[1].data.numpy() |
98 | 98 | accuracy = float((pred_y == test_y).astype(int).sum()) / float(test_y.size) |
99 | 99 | print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.numpy(), '| test accuracy: %.2f' % accuracy) |
100 | 100 |
|
101 | 101 | # print 10 predictions from test data |
102 | 102 | test_output = rnn(test_x[:10].view(-1, 28, 28)) |
103 | | -pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze() |
| 103 | +pred_y = torch.max(test_output, 1)[1].data.numpy() |
104 | 104 | print(pred_y, 'prediction number') |
105 | 105 | print(test_y[:10], 'real number') |
106 | 106 |
|
0 commit comments