skorchでのエラーについて

Question

skorchを使ってMNISTを使いたいと思い、下記のようなコードを実行しましたがエラーが出てしまいました。解決策はありますでしょうか。試したこととしてnp.expand_dims(x, 1)をやってみましたが結果は変わりませんでした。

%matplotlib inline
from skorch import NeuralNet
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
torch.manual_seed(123)

data_x, data_y = fetch_openml('mnist_784', data_home = './data', cache=True, return_X_y=True)
x_train, x_test, y_train, y_test = train_test_split(data_x, data_y, test_size=0.25, random_state=42)

x_train = np.expand_dims(x_train / 255, 1).astype(np.float32)
x_test = np.expand_dims(x_test / 255, 1).astype(np.float32)

y_train = y_train.astype(np.int64)
y_test = y_train.astype(np.int64)

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(1600, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))    # (32, 13, 13)
        x = F.relu(F.max_pool2d(self.conv2_drop(
            self.conv2(x)), 2))   # (64, 5, 5)
        x = x.view(-1, x.size(1) * x.size(2) * x.size(3))
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        x = F.log_softmax(x, dim=-1)
        return x

device = torch.device('cuda')

model = NeuralNet(
  CNN,            # ここで、先程定義したNetクラスを引数として与える
  max_epochs=10,
  optimizer=torch.optim.Adam,
  lr=0.001,
  device=device,
  batch_size=128,
  train_split=None,
  criterion=nn.NLLLoss    # CNNの最後のactivationがlog_softmaxなので、lossはNLLoss。
)


model.fit(X_train, y_train)

pred = model.prefict(x_test)  

pred = pred.argmax(axis=1)
print("acc:{}".format(accuracy_score(y_test, pred)))

RuntimeError: Expected 4-dimensional input for 4-dimensional weight [32, 1, 3, 3], but got 2-dimensional input of size [128, 784] instead

Answer 1

そもそも配列の形状がおかしかった