pytorchで混同行列

前提・実現したいこと

pytorchで混同行列を用い識別率を計測したいです。以下に示しているプログラムはMNISTデータを１０クラスに分類するものです。
自分がやることは10×10の行列を作り、行列[推論結果][正解ラベル]に1を足すことだと考えています。

実現できないところ（自分がわからないところ）

このネットワークでテストデータを識別した際に、どのクラス（0~9）に識別されたかを確認することです。
例えば、4の手書き数字画像を入力すると「4」と識別されるとして、この出力の4を知りたいです。

該当のソースコード

python
1#学習用
2mport os
3import sys
4import numpy as np
5import datetime
6
7import torch
8import torch.nn as nn
9import torch.nn.functional as F
10import torch.optim as optim
11
12import getdata
13
14PATH_MNIST = '/mnist'
15PATH_RESULT = 'result'
16
17### definition of the network
18#
19class NN(nn.Module):
20
21    def __init__(self):
22        super(NN, self).__init__()
23        Hin1, Win1 = 28, 28
24        self.conv1 = nn.Conv2d(1, 32, kernel_size = 5)
25        Hout1, Wout1 = _output_shape(Hin1, Win1, self.conv1) # 24 x 24
26        self.conv2 = nn.Conv2d(self.conv1.out_channels, 64, kernel_size = 5) 
27        Hout2, Wout2 = _output_shape(Hout1//2, Wout1//2, self.conv2) # 8 x 8
28        self.fc1 = nn.Linear(64 * (Hout2//2) * (Wout2//2), 1024) # 64 x 4 x 4
29        self.fc2 = nn.Linear(self.fc1.out_features, 10)
30
31    def forward(self, X):
32        X = F.relu(F.max_pool2d(self.conv1(X), 2))
33        X = F.relu(F.max_pool2d(self.conv2(X), 2))
34        X = X.view(-1, self.fc1.in_features)
35        X = F.relu(self.fc1(X))
36        X = self.fc2(X)
37        return F.log_softmax(X, dim = 1)
38
39def _output_shape(Hin, Win, conv2d):
40        Hout = int(np.floor((Hin + 2 * conv2d.padding[0] - conv2d.dilation[0] * (conv2d.kernel_size[0] - 1) - 1) / conv2d.stride[0] + 1))
41        Wout = int(np.floor((Win + 2 * conv2d.padding[1] - conv2d.dilation[1] * (conv2d.kernel_size[1] - 1) - 1) / conv2d.stride[1] + 1))
42        return Hout, Wout
43
44def evaluate(model, X, Y, bindex):
45
46    nbatch = bindex.shape[0]
47    loss = 0
48    ncorrect = 0
49    with torch.no_grad():
50        for ib in range(nbatch):
51            ii = np.where(bindex[ib, :])[0]
52            output = model(X[ii, ::])
53            #loss += F.nll_loss(output, Y[ii], size_average=False).item()
54            loss += F.nll_loss(output, Y[ii], reduction='sum').item()
55            pred = output.max(1, keepdim=True)[1] # argmax of the output
56            ncorrect += pred.eq(Y[ii].view_as(pred)).sum().item()
57
58    loss /= X.shape[0]
59    acc = ncorrect / X.shape[0]
60
61    return loss, acc
62
63    
64
65if __name__ == '__main__':
66
67    ### device
68    #
69    use_gpu_if_available = True
70    if use_gpu_if_available and torch.cuda.is_available():
71        device = torch.device('cuda')
72    else:
73        device = torch.device('cpu')
74    print('# using', device)
75    
76    ### reading and preparing the training data
77    #
78    data = getdata.Data(PATH_MNIST, nV = 10000)
79
80    D = data.nrow * data.ncol
81    K = data.nclass
82    datLraw, labL = data.getData('L')
83    datL = datLraw.reshape((-1, 1, data.nrow, data.ncol))
84    datVraw, labV = data.getData('V')
85    datV = datVraw.reshape((-1, 1, data.nrow, data.ncol))
86    NV = datV.shape[0]
87    NL = datL.shape[0]
88
89    ### to torch.Tensor
90    #
91    XL = torch.from_numpy(datL.astype(np.float32)).to(device)
92    YL = torch.from_numpy(labL).to(device)
93    XV = torch.from_numpy(datV.astype(np.float32)).to(device)
94    YV = torch.from_numpy(labV).to(device)
95    
96    ### initializing the network
97    #
98    torch.manual_seed(0)
99    nn = NN()
100    model = nn.to(device)
101    print(nn)
102    optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum = 0.9)
103    print(optimizer)
104
105    ### training
106    #
107    batchsize = 100
108    bindexL = getdata.makeBatchIndex(NL, batchsize)
109    nbatchL = bindexL.shape[0]
110    bindexV = getdata.makeBatchIndex(NV, batchsize)
111    nbatchV = bindexV.shape[0]
112
113    nitr = 2000
114    nd = 0
115    start = datetime.datetime.now()
116
117    for i in range(nitr):
118
119        if (i < 500 and i % 100 == 0) or (i % 500 == 0):
120            
121            model.eval()  # setting the module in evaluation mode
122            lossL, accL = evaluate(model, XL, YL, bindexL)
123            lossV, accV = evaluate(model, XV, YV, bindexV)
124            print('{} {:.2f}'.format(i, nd/NL), end = '     ')
125            print('{:.4f} {:.2f}'.format(lossL, accL*100), end = '     ') 
126            print('{:.4f} {:.2f}'.format(lossV, accV*100), end = '     ') 
127            print()
128            
129        model.train()  # setting the module in training mode
130        ib = np.random.randint(0, nbatchL)
131        ii = np.where(bindexL[ib, :])[0]
132        optimizer.zero_grad()
133        output = model(XL[ii, :])
134        loss = F.nll_loss(output, YL[ii])
135        loss.backward()
136        optimizer.step()
137
138        nd += ii.shape[0]
139
140    print('# elapsed time: ', datetime.datetime.now() - start)
141    
142    ### saving the model
143    #
144    fnModel = os.path.join(PATH_RESULT, os.path.splitext(sys.argv[0])[0] + '-params.pickle')
145    with open(fnModel, mode = 'wb') as f:
146        torch.save(model.state_dict(), f)
147        print('# The model is saved to ', fnModel)

python
1#推論用
2import os
3import sys
4import pickle
5import numpy as np
6import datetime
7
8import torch
9import torch.nn as nn
10import torch.nn.functional as F
11
12import getdata
13import cnnL
14
15
16if __name__ == '__main__':
17
18    ### device
19    #
20    use_gpu_if_available = True
21    if use_gpu_if_available and torch.cuda.is_available():
22        device = torch.device('cuda')
23    else:
24        device = torch.device('cpu')
25    print('# using', device)
26
27    ### initializing the network
28    #
29    fnModel = os.path.join(cnnL.PATH_RESULT,'cnnL-params.pickle')
30    torch.manual_seed(0)
31    nn = cnnL.NN()
32    with open(fnModel, mode = 'rb') as f:
33        nn.load_state_dict(torch.load(f))
34    model = nn.to(device)
35
36    ### reading and preparing the training data
37    #
38    data = getdata.Data(cnnL.PATH_MNIST, nV = 10000)
39
40    D = data.nrow * data.ncol
41    K = data.nclass
42    datTraw, labT = data.getData('T')
43    datT = datTraw.reshape((-1, 1, data.nrow, data.ncol))
44    NT = datT.shape[0]
45
46    ### to torch.Tensor
47    #
48    XT = torch.from_numpy(datT.astype(np.float32)).to(device)
49    YT = torch.from_numpy(labT).to(device)
50    
51    ### evaluation
52    #
53    batchsize = 100
54    bindexT = getdata.makeBatchIndex(NT, batchsize)
55    model.eval()  # setting the module in evaluation mode
56    start = datetime.datetime.now()
57    lossT, accT = cnnL.evaluate(model, XT, YT, bindexT)
58    print('# elapsed time: ', datetime.datetime.now() - start)
59    print('{:.4f} {:.2f}'.format(lossT, accT*100))
60            
61