Как преобразовать тензор 3D-изображения в формат PyTorch MNIST для вывода?

Я обучил следующую модель с набором данных Pytorch MNIST и сохранил ее в конце.

Теперь я хочу сделать выводы на некоторых изображениях, которые не являются частью набора данных MNIST. Изображения имеют формат jpg и загружаются с помощью OpenCV как (28, 28) массивы NumPy (см. встроенные комментарии для типов и форм):

При выполнении вывода модель всегда предсказывает "3" для любого тестового образца, и я получаю следующее предупреждение:

Похоже, я неправильно отформатировал свои изображения для вывода. Я попытался получить доступ к изображениям в загрузчике данных обучающего файла, но в итоге получил объект DataLoader, неспособный получить доступ к формам изображений. Можете ли вы сказать мне, как мне нужно отформатировать свои изображения для вывода?

from torchvision import datasets
from torchvision.transforms import ToTensor
from torch.utils.data import DataLoader

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim


train_data = datasets.MNIST(
    root="data",
    train=True,
    transform=ToTensor(),
    download=True
    )

test_data = datasets.MNIST(
    root="data",
    train=False,
    transform=ToTensor(),
    download=True
    )


loaders = {
    "train": DataLoader(train_data,
                        batch_size=100,
                        shuffle=True,
                        num_workers=0),

    "test": DataLoader(test_data,
                    batch_size=100,
                    shuffle=True,
                    num_workers=0)
}

class CNN(nn.Module):

    def __init__(self):
        super(CNN, self).__init__()

        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)

        return F.softmax(x)
    

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = CNN().to(device)

optimizer = optim.Adam(model.parameters(),
                    lr=0.001)

loss_fn = nn.CrossEntropyLoss()


def train(epoch):

    model.train()
    for batch_idx, (data, target) in enumerate(loaders["train"]):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 20 == 0:
            print(f'Train epoch: {epoch} [{batch_idx * len(data)}/{len(loaders["train"].dataset)} ({100. * batch_idx / len(loaders["train"]):.0f}%)]\t{loss.item():.6f}')


def test():
    model.eval()

    test_loss = 0
    correct = 0

    with torch.no_grad():
        for data, target in loaders["test"]:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += loss_fn(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(loaders['test'].dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(loaders["test"].dataset)} ({100. * correct / len(loaders["test"].dataset):.0f}%\n)')


for epoch in range(1, 11):
    train(epoch)
    test()

torch.save(model.state_dict(), "model.pth")

import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.transforms import Compose, ToTensor


# same class as in the file above
class CNN(nn.Module):

        def __init__(self):
            super(CNN, self).__init__()

            self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
            self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
            self.conv2_drop = nn.Dropout2d()
            self.fc1 = nn.Linear(320, 50)
            self.fc2 = nn.Linear(50, 10)

        def forward(self, x):
            x = F.relu(F.max_pool2d(self.conv1(x), 2))
            x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
            x = x.view(-1, 320)
            x = F.relu(self.fc1(x))
            x = F.dropout(x, training=self.training)
            x = self.fc2(x)

            return F.softmax(x)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = CNN().to(device)

model.load_state_dict(torch.load("model.pth"))


img = cv2.imread("image.jpg", cv2.IMREAD_GRAYSCALE)

print(type(img))  # Output: <class 'numpy.ndarray'>
print(img.shape)  # Output: (28, 28)

model.eval()

with torch.no_grad():

    transform = Compose([ToTensor()])
    img = transform(img)

    print(type(img))  # Output: <class 'torch.Tensor'>
    print(img.shape)  # Output: torch.Size([1, 28, 28])

    output = model(img)
    pred = output.argmax(dim=1, keepdim=True)
    print(pred)

UserWarning: dropout2d: Received a 3D input to dropout2d and assuming that channel-wise 1D dropout behavior is desired - input is interpreted as shape (N, C, L), where C is the channel dim. This behavior will change in a future release to interpret the input as one without a batch dimension, i.e., shape (C, H, W). To maintain the 1D channel-wise dropout behavior, please switch to using dropout1d instead.
  warnings.warn("dropout2d: Received a 3D input to dropout2d and assuming that channel-wise "
c:/Users/u/project/venv/src/digits_model_prediction.py:28: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.
  return F.softmax(x)