Pytorch学习笔记（3）:前馈神经网络

最新推荐文章于 2024-01-17 10:36:23 发布

马可仕马可仕

最新推荐文章于 2024-01-17 10:36:23 发布

阅读量2.1k

点赞数 1

分类专栏： Pytorch在学了文章标签：神经网络 python pytorch

本文链接：https://blog.csdn.net/weixin_43554584/article/details/111500962

版权

本博客基于github上某位大神的pytorch入门学习代码，在他的基础上加上了更详细的中文注释以及不懂的模块使用方法。github连接：https://github.com/yunjey/pytorch-tutorial

前馈神经网络

在运行代码之前，请确保当前环境下已安装torch、torchvision。

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# 配置代码运行所用的设备
# 如果有cuda就使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 超参数设置
input_size = 784
hidden_size = 500
num_classes = 10
num_epoch = 5
batch_size = 100
learning_rate = 0.001

# 加载训练集和测试集
train_dataset = torchvision.datasets.MNIST(root="../../data/minist", train=True, transform=transforms.ToTensor(), download=True)
test_dataset = torchvision.datasets.MNIST(root="../../data/minist", train=False, transform=transforms.ToTensor()<