【PyTorch】分類結果を表で可視化する – confusion_matrixの作成

Python

confusion_matrix

PyTorchでニューラルネットワークを作成、画像分類したときの結果を表にする方法です。

例えばCIFA10を認識させたときに、ネコの画像を正しくネコと認識できた回数、間違えてイヌと認識してしまった回数…を画面上に表示させます。

ネコの画像はイヌの画像に似てるから誤認識しやすい、逆にネコは飛行機に似てないから飛行機とは認識されづらい、といった分析が可能になります。

PyTorchで実装

マトリックスを作成する関数をつくってみました。

def make_matrix(matrix,predicted,target):
    p = predicted.to('cpu').numpy().T
    t = target.data.view_as(predicted).to('cpu').numpy().T
    for i in range(len(p)):
        for j in range(len(p[0])):
                matrix[p[i][j]][t[i][j]]+=1

これを学習後にmain.pyのどこかで呼び出すとマトリックスがターミナル上に表示されます。

main.py内のtest関数で実際にマトリックスを作成するときのpythonコードです。

CIFAR10は10クラスの分類なので10×10の表が作成されます。

def test(args,model,test_loader,criterion):
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    matrix=np.zeros((10,10))
    for data, target in test_loader:
        
        data, target = data.to('cuda:0'), target.to('cuda:0')
        output = model(data)
        test_loss += criterion(output, target).item()

        predicted = output.max(1,keepdim=True)[1]
        total += target.size(0)
        correct += predicted.eq(target.data.view_as(predicted)).to('cpu').sum()
        make_matrix(matrix,predicted,target)
      
    acc = 100. * correct/len(test_loader.dataset)

    test_loss /= len(test_loader)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

    print(matrix)

    return acc

コメント

タイトルとURLをコピーしました