PyTorch中怎么加载数据集

1949idc 2年前 (2024-09-18) 阅读数 318 #深度学习

在 PyTorch 中,可以使用 torchvision.datasets 模块来加载常见的数据集,如 MNIST、CIFAR-10 等。这些数据集通常会被下载到本地,并返回一个 Dataset 对象,可以通过 DataLoader 对象来对数据集进行批量加载和随机打乱。

以下是一个加载 MNIST 数据集的示例代码:

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载 MNIST 训练集和测试集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 创建 DataLoader 对象
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 遍历数据集
for inputs, labels in train_loader:
    # 在这里进行模型训练
    pass

上面的代码首先定义了数据预处理的方法 transform,然后使用 datasets.MNIST 加载了 MNIST 数据集的训练集和测试集,并创建了对应的 DataLoader 对象 train_loadertest_loader。最后,可以通过遍历 train_loader 来逐批获取训练数据和标签,并进行模型训练。

版权声明

本文内容由互联网用户自发贡献,该文观点仅代表作者本人
本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。

© 2010 首途云安 & 厦门硕顿信息技术有限公司 & 闽ICP备11016866号  增值电信业务经营许可证:B1-20203020 地址:福建厦门思明区嘉禾路297号1806
高新技术企业
软件产品证书
计算机软件著作权
ISO认证
国家3A企业