Как создать набор данных Pytorch из файлов .pt?

Question

Я преобразовал изображения MNIST, сохраненные в виде файлов .pt в папку на диске Google. Я пишу свой код Pytorch в Colab.

Я хотел бы использовать эти файлы и создать набор данных, в котором эти изображения хранятся в виде тензоров. Как я могу это сделать?

Преобразование изображений во время тренировки заняло слишком много времени. Следовательно, преобразовал их и сохранил их как файлы .pt. Я просто хочу загрузить их обратно как набор данных и использовать их в моей модели.

Answer 1

Подход, который вы используете для сохранения изображений, действительно хорошая идея.В таком случае вы можете просто написать свой собственный класс набора данных для загрузки изображений.

from torch.utils.data import Dataset, DataLoader
from torch.utils.data.sampler import RandomSampler

class ReaderDataset(Dataset):
    def __init__(self, filename):
        # load the images from file

    def __len__(self):
        # return total dataset size

    def __getitem__(self, index):
        # write your code to return each batch element

Затем вы можете создать Dataloader следующим образом.

train_dataset = ReaderDataset(filepath)
train_sampler = RandomSampler(train_dataset)
train_loader = DataLoader(
    train_dataset,
    batch_size=args.batch_size,
    sampler=train_sampler,
    num_workers=args.data_workers,
    collate_fn=batchify,
    pin_memory=args.cuda,
    drop_last=args.parallel
)
# args is a dictionary containing parameters
# batchify is a custom function that prepares each mini-batch