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

最新推荐文章于 2024-03-09 10:13:03 发布

原创最新推荐文章于 2024-03-09 10:13:03 发布

· 280 阅读

1 ·

版权

文章标签：

#pytorch #深度学习 #python

机器学习专栏收录该内容

8 篇文章

订阅专栏

本文介绍了一个自定义PyTorch数据集类的具体实现，该类能够从指定目录加载图像，并利用TensorBoard进行数据可视化展示。通过遍历数据集中的每一项并将其转换为NumPy数组形式，最终使用TensorBoard将图像数据可视化。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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")