CIFAR10のチュートリアルをpytorchでやってみてそれを発展させてクラスごとのconfusion

やりたいこと

cifar10のdataseを使ってチュートリアルをもとに精度クラスごとに出すことに成功しました。

それでいて、次にやろうと思っているのは、その各クラスが本当にあっているかを確かめたくてconfusion_matrixでの比較を行いたいと考えています。

いくつか見ていると、sklearnのconfusion_matrics（正解ラベル,予測ラベル）を使うとできるらしいということがわかりました。引数は正解と予測のラベルだということがわかりました。

しかし、実際にpytorchで出した、正解ラベル（list）と予測ラベル(list)をどのようにして出したらいいのかわかりません。

そのほか、以下codeのsklearnのようにデータセットをx,yでスプリットできたらいいのですが,,,やり方がわかりません。

python
1from sklearn import svm, datasets
2from sklearn.model_selection import train_test_split
3from sklearn.metrics import plot_confusion_matrix
4
5iris = datasets.load_iris()
6X = iris.data
7y = iris.target
8class_names = iris.target_names
9
10# Split the data into a training set and a test set
11X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

以下に理想図をのせます（こんな感じになってほしい）

array([[806,  28,  52,  12,  16,   1,   9,   7,  43,  26],
       [ 47, 832,   4,   4,   1,   0,  18,   1,  15,  78],
       [ 95,   8, 474,  47, 146,  33, 131,  34,  10,  22],
       [ 31,  18,  55, 478,  85,  90, 144,  45,  19,  35],
       [ 43,   4,  37,  37, 686,   5, 115,  64,   5,   4],
       [ 24,   6,  36, 181,  92, 443,  96,  87,  12,  23],
       [ 12,   3,  28,  39,  25,   6, 873,   4,   6,   4],
       [ 45,  12,  18,  30,  87,  20,  23, 732,   5,  28],
       [137,  43,   9,   9,  13,   0,   8,   2, 752,  27],
       [ 81, 164,   5,  11,   4,   1,  12,  11,  16, 695]])

実際のコード＋結果

長いので、該当箇所を貼っていきます。（もし、このほかに需要な箇所があればご指摘くだされば添付いたします。）

class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        images, labels = images.to(device), labels.to(device) 
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        # print(classes)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            # print(label)
            class_correct[label] += c[i].item()
            # print(class_correct[label])
            class_total[label] += 1


for i in range(10):
    print('Accuracy of %d %5s : %2d %%' % (i,classes[i], 100 * class_correct[i] / class_total[i]))

from sklearn.metrics import confusion_matrix
from sklearn.metrics import plot_confusion_matrix
cm = confusion_matrix(class_correct, class_total)
print(cm)
disp = plot_confusion_matrix(cm)
plt.show()

実行結果

Accuracy of 0 plane : 71 %
Accuracy of 1   car : 77 %
Accuracy of 2  bird : 29 %
Accuracy of 3   cat : 57 %
Accuracy of 4  deer : 49 %
Accuracy of 5   dog : 30 %
Accuracy of 6  frog : 62 %
Accuracy of 7 horse : 69 %
Accuracy of 8  ship : 72 %
Accuracy of 9 truck : 72 %
[[0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 1]
 [0 0 0 0 0 0 0 0 0 0 0]]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-24-bfa9f43d16cb> in <module>()
     24 cm = confusion_matrix(class_correct, class_total)
     25 print(cm)
---> 26 disp = plot_confusion_matrix(cm)
     27 plt.show()
     28 

TypeError: plot_confusion_matrix() missing 2 required positional arguments: 'X' and 'y_true'

最後に

以上のことから、matricsの図を描くためのヒントのようなものでもいただければ助かります。
特に自分のコードにおける予測ラベルと、正解ラベルのセットがどこにあるかが認識できればmatrixがかけるのでは？とかってに思っています。
また、x,yの座標が取得できればそれをもとに図としておこすことができるのではないかと思っています。

ですので、アドバイスを頂ければと思います。

どうかよろしくお願いします