>- ** 本文为[365天深度学习训练营]中的学习记录博客**
>- ** 原作者:[K同学啊]**
我的环境:
如果设备上支持GPU就使用GPU,否则使用CPU
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision
# 设置硬件设备,如果有GPU则使用,没有则使用cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device
device(type='cuda') # 使用的是GPU
使用dataset下载MNIST数据集,并划分好训练集与测试集
使用dataloader加载数据,并设置好基本的batch_size
⭐ torchvision.datasets.MNIST详解
torchvision.datasets
是Pytorch自带的一个数据库,我们可以通过代码在线下载数据,这里使用的是torchvision.datasets
中的MNIST
数据集。
函数原型:
torchvision.datasets.MNIST(root, train=True, transform=None, target_transform=None, download=False)
参数说明:
True
-训练集,False
-测试集True
,从互联网上下载数据集,并把数据集放在root目录下。train_ds = torchvision.datasets.MNIST('data',
train=True,
transform=torchvision.transforms.ToTensor(), # 将数据类型转化为Tensor
download=True)
test_ds = torchvision.datasets.MNIST('data',
train=False,
transform=torchvision.transforms.ToTensor(), # 将数据类型转化为Tensor
download=True)
⭐ torch.utils.data.DataLoader详解
torch.utils.data.DataLoader
是Pytorch自带的一个数据加载器,结合了数据集和取样器,并且可以提供多个线程处理数据集。
函数原型:
torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=None, sampler=None, batch_sampler=None, num_workers=0, collate_fn=None, pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None, multiprocessing_context=None, generator=None, *, prefetch_factor=2, persistent_workers=False, pin_memory_device='')
参数说明: