import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Hyper-parameters
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001
# MNIST dataset
train_dataset = torchvision.datasets.MNIST(root='../../data',
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = torchvision.datasets.MNIST(root='../../data',
train=False,
transform=transforms.ToTensor())
# Data loader
train_loader = torch
Pytorch入门 1.4 feedforward neural network
最新推荐文章于 2024-02-16 22:15:53 发布
这篇博客记录了使用Pytorch构建并训练一个具有单层隐藏层的feedforward神经网络的过程。在5个epoch的训练中,损失逐渐降低,最终在测试集上达到了97.83%的高准确率,利用了cuda进行运算加速。
摘要由CSDN通过智能技术生成