深度学习-PyTorch的批数据训练

最新推荐文章于 2021-07-05 10:46:59 发布

原创

最新推荐文章于 2021-07-05 10:46:59 发布 · 970 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #pytorch #批数据训练

本文介绍在处理大量数据时，如何利用PyTorch的TensorDataset和DataLoader进行批数据训练，通过设置批次大小有效地进行模型训练。

当数据比较大时，PyTorch提供一个工具可以将数据分批进行训练。思想就是将数据放入torch中的TensorDataset数据库，再放入数据加载器，设定一批的加载个数进行加载。

官方文档解释：

import torch
import torch.utils.data as Data

BATCH_SIZE = 5  # 一批5个数据
x = torch.linspace(1, 10, 10)  # 1，2，3，4，5，6，7，8，9，10
y = torch.linspace(10, 1, 10)  # 10，9，8，7，6，5，4，3，2，1
torch_dataset = Data.TensorDataset(x, y)  # 定义数据库
loader = Data.DataLoader(  # 数据加载器
    dataset=torch_dataset,
    batch_size=BATCH_SIZE,  # 设置每个数据段大小
    shuffle=True,  # 是否打乱顺序
)

for epoch in range(3):  # 循环执行三次
    for step, (batch_x, batch_y) in enumerate(loader):  # 每个循环中10个数据分成两次训练
        print("epoch:", epoch, "|step:", step, "|batch x:", batch_x.numpy(), "|batch y:", batch_y.numpy())