Pytorch学习笔记【18】：使用GPU加速

最新推荐文章于 2024-12-26 11:24:05 发布

strong tyj

最新推荐文章于 2024-12-26 11:24:05 发布

阅读量1k

点赞数 1

分类专栏： # Pytorch

本文链接：https://blog.csdn.net/qq_36499794/article/details/103178278

版权

本文基于之前CNN识别手写数字的博客，介绍如何利用GPU加速Pytorch模型训练，包括改过的代码、原本代码和改过后的对比。

摘要由CSDN通过智能技术生成

之前我写过一篇CNN识别手写数字的博客，我这一篇的介绍将基于那一篇的代码做出相关改进

1. 改过的代码

2. 原本的代码

import os

# third-party library
import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt


# 定义一些参数
EPOCH = 1               # 训练数据的次数，我们这里假定训练一次
BATCH_SIZE = 50         # 每次训练的数据量，这个会产生每一次训练分多少次进行，或者多少批进行
LR = 0.001              # 学习率
DOWNLOAD_MNIST = False


# 下载并且加载数据集
if not(os.path.exists('./mnist/')) or not os.listdir('./mnist/'):
    # not mnist dir or mnist is empyt dir
    DOWNLOAD_MNIST = True

train_data = torchvision.datasets.MNIST(
    root='./mnist/',
    train=True,                                     # 表示训练数据
    transform=torchvision.transforms.ToTensor(),    # 将数据转换成tensor
                                                    # torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]
    download=DOWNLOAD_MNIST,
)

# plot one example
print(train_data.train_data.size())                 # (60000, 28, 28)
print(train_data.train_labels.size())