PyTorch 数据集使用

Dataset

import torchvision
from torch.utils.tensorboard import SummaryWriter

dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])
train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=dataset_transform, download=True)
test_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=dataset_transform, download=True)

# print(test_set[0])
# print(test_set.classes)

writer = SummaryWriter('P11')
for i in range(10):
    img, target = test_set[i]
    writer.add_image("test_set", img, i)

writer.close()

Dataloader 

import torchvision
from torch.utils.data import DataLoader

# 准备的测试数据集
test_data = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor())

test_loader = DataLoader(dataset=test_data, batch_size=4, shuffle=False, num_workers=0, drop_last=False)

#测试集中第一张图片及target
img, target = test_data[0]
print(img.shape)
print(target)

for data in test_loader:
    imgs, targets = data
    print(imgs.shape)
    print(targets)

第一张的print结果

torch.Size([3, 32, 32])
3

 [3, 32, 32] -- 3通道，32*32 大小

3 -- 第三类

Batch size = 4时的print结果

torch.Size([4, 3, 32, 32])
tensor([3, 8, 8, 0])
torch.Size([4, 3, 32, 32])
tensor([6, 6, 1, 6])
torch.Size([4, 3, 32, 32])
tensor([3, 1, 0, 9])
...
...
...

 [4, 3, 32, 32] -  4张3通道的32*32图片

[3, 8, 8, 0] - 4张图片分别是 第3，8，8，0类

用tensorboard查看load结果

from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

test_data = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor())

test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=False, num_workers=0, drop_last=False)

#测试集中第一张图片
img, target = test_data[0]
print(img.shape)
print(target)

writer = SummaryWriter('dataloader')
step = 0
for data in test_loader:
    imgs, targets = data
    writer.add_images("test_data", imgs, step)
    step = step+1

writer.close()

⚠️： 最后一组数据不足64张，若想要丢弃改动如下:

1. test_loader里 drop_last = True

2. writer.add_images里 test_data_drop_last

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

test_data = torchvision.datasets.CIFAR10("./dataset", train=False, transform=torchvision.transforms.ToTensor())

test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=False, num_workers=0, drop_last=True)

#测试集中第一张图片
img, target = test_data[0]
print(img.shape)
print(target)

writer = SummaryWriter('dataloader')
step = 0
for data in test_loader:
    imgs, targets = data
    writer.add_images("test_data_drop_last", imgs, step)
    step = step+1

writer.close()