Как в датасет MNIST, положить дополнительно свой с добавлением классов?

Question

[venom99]

Нашёл почти готовый проект на github, не могу понять как адаптировать его под свою задачу. Есть датасет MNIST, в нём классы от 0 - 9, сделал датасет с картинками знака бесконечности 28х28, произвёл конкатенацию, и не понимаю как научить нейросеть понимать, что это знак бесконечности, а не число 8.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
import cv2
import os
import pandas as pd
import torch.optim as optim
import torch.utils as utils
from torchvision import transforms, datasets

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = cv2.imread(img_path, 0)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label


device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
batch_size = 400
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_data = datasets.MNIST('./MNIST-main/data', True, transform=transforms.ToTensor(), download=False)
train_dir = "./MNIST-main/data/lrTrain"
train_datasets = CustomImageDataset('MNIST-main/data/labales.csv', 'MNIST-main/data/lrTrain/awd',transform=transforms.ToTensor())
composed = torch.utils.data.ConcatDataset([train_data, train_datasets])

print(train_data.__len__())
train_loader = utils.data.DataLoader(train_data, batch_size=batch_size,
                                     shuffle=True, num_workers=2)

test_data = datasets.MNIST('./MNIST-main/data', False, transform=transforms.ToTensor(), download=False)
composed1 = torch.utils.data.ConcatDataset([test_data, train_datasets])
test_loader = utils.data.DataLoader(test_data, batch_size=batch_size,
                                    shuffle=True, num_workers=2)


class net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(256, 120)
        self.fc2 = nn.Linear(120, 60)
        self.fc3 = nn.Linear(60, 10)
        self.pool = nn.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 256)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return F.log_softmax(x, -1)


if __name__ == '__main__':
    model = net().to(device)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    epochs = 30
    for i in model._get_name():
        print(i)
        print(']]]]]]]]]]]]]]]')
    model.load_state_dict(torch.load('MNIST-main/MNIST.pth'))
    model.eval()
    with torch.inference_mode():
        # for img, label in test_loader:
        #     img, label = img.to(device), label.to(device)
        #     out = model(img)
        #     print(out)
        #     loss = loss_fn(out, label)
        img, label = next(iter(test_loader))
        img, label = img.to(device), label.to(device)

        outputs = model(img.to(device))
        _, preds = torch.max(outputs, 1)
        preds = preds.cpu().numpy()
        classes = label.cpu().numpy()
        print(preds)
        print(classes)
        import matplotlib.pyplot as plt

        plt.imshow(img[0].cpu().squeeze(0), cmap='gray')
        plt.show()

Если в train_loader и test_loader передать train_datasets, то есть полностью обойтись без MNIST, то в print(preds) и print(classes) получаю следующее.
[8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8]
[11 11 11 11 11 11 11 11 11 11 11 11 11 11 11 11 11 11 11]
В labales.csv всем картинкам задаю метку 11, а вот что делать со значением 8, не знаю.