使用Dataset读取数据，并用tensorboard显示

from torch.utils.data import Dataset
from PIL import Image
import os
from torch.utils.tensorboard import SummaryWriter
import numpy as np


class MyData(Dataset):

    def __init__(self, root_dir):
        self.root_dir = root_dir
        # self.label_dir = label_dir
        # self.path = os.path.join(self.root_dir, self.label_dir)
        self.img_path = os.listdir(self.root_dir)

    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir, img_name)
        img = Image.open(img_item_path)
        return img

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

root_dir = "mnist_data" # root_dir 是 数据集的相对路径
dataset = MyData(root_dir)

writer = SummaryWriter("logs")
for i in range(len(dataset)):
    img_array = np.array(dataset[i])
    writer.add_image("mnist", img_array, i, dataformats='HW')

writer.close()
print("end")

结果如下：