Pytorch学习笔记3.数据集、dataloader

数据集的使用

import torchvision
from torch.utils.tensorboard import SummaryWriter

dataset_transform=torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])
//这里使用CIFAR10包
train_set=torchvision.datasets.CIFAR10(root="./dataset",train=True,transform=dataset_transform,download=True)//训练集
test_set=torchvision.datasets.CIFAR10(root="./dataset",train=False,transform=dataset_transform,download=True)//测试集

writer=SummaryWriter("p10")
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
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=True,num_workers=0,drop_last=False)
# batch_size=4： 每次从数据集中抓4个
# drop_last 抓取剩余的，舍去或保留

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

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

writer.close()