pytorch神经网络入门代码

hehui0921

已于 2024-02-16 10:38:13 修改

阅读量683

点赞数 2

分类专栏：神经网络文章标签： pytorch 神经网络人工智能

于 2024-02-16 09:51:12 首次发布

本文链接：https://blog.csdn.net/oHeHui1/article/details/136126207

版权

本文介绍了如何使用PyTorch库在MNIST数据集上构建和训练两种类型的神经网络：简单的全连接网络和卷积神经网络(CNN)。代码展示了从数据预处理、模型定义、训练过程到测试的完整流程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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

# 定义神经网络结构
class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

# 设置超参数
input_size = 784  # MNIST数据集的输入大小是28x28=784
hidden_size = 784
num_classes = 10


learning_rate