基于PyTorch,如何构建一个简单的神经网络

建立神经网络
import osimport torchfrom torch import nnfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transforms
加载训练设备
device = 'cuda' if torch.cuda.is_available() else 'cpu' #检测gpu是否可用,不可用使用cpuprint('Using {} device'.format(device)) #输出使用设备类型
定义类
class NeuralNetwork(nn.Module):? ?def __init__(self): #定义网络结构? ? ? ?super(NeuralNetwork, self).__init__()? ? ? ?self.flatten = nn.Flatten()? ? ? ?self.linear_relu_stack = nn.Sequential(? ? ? ? ? ?nn.Linear(28*28, 512),? ? ? ? ? ?nn.ReLU(),? ? ? ? ? ?nn.Linear(512, 512),? ? ? ? ? ?nn.ReLU(),? ? ? ? ? ?nn.Linear(512, 10),? ? ? ? ? ?nn.ReLU()? ? ? ?)? ?def forward(self, x): #前向传播? ? ? ?x = self.flatten(x)? ? ? ?logits = self.linear_relu_stack(x)? ? ? ?return logits
model = NeuralNetwork().to(device) #实例化模型print(model)
X = torch.rand(1, 28, 28, device=device) ?#生成(1,28,28)的数据logits = model(X) #向模型输入数据pred_probab = nn.Softmax(dim=1)(logits) #调用softmax 将预测值映射为(0,1)间的概率y_pred = pred_probab.argmax(1) #最大概率对应分类print(f"Predicted class: {y_pred}")
神经网络各层说明
input_image = torch.rand(3,28,28) #生成(3,28,28)的数据print(input_image.size())
nn.Flatten 层
flatten = nn.Flatten()flat_image = flatten(input_image) #(3,28,28)转换为(3,784)print(flat_image.size())
nn.Linear 层
layer1 = nn.Linear(in_features=28*28, out_features=20) #输入(3,28*28) 输出(3,20)hidden1 = layer1(flat_image)print(hidden1.size())
nn.ReLU 层
print(f"Before ReLU: {hidden1}\n\n")hidden1 = nn.ReLU()(hidden1)print(f"After ReLU: {hidden1}")
nn.Sequential 层
softmax = nn.Softmax(dim=1)pred_probab = softmax(logits)
输出模型结构
print("Model structure: ", model, "\n\n")for name, param in model.named_parameters():? ?print(f"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \n")

关注公众号:拾黑(shiheibook)了解更多
[广告]赞助链接:
四季很好,只要有你,文娱排行榜:https://www.yaopaiming.com/
让资讯触达的更精准有趣:https://www.0xu.cn/
关注网络尖刀微信公众号随时掌握互联网精彩
赞助链接
排名
热点
搜索指数
- 1 习近平将发表二〇二六年新年贺词 7904141
- 2 2026年国补政策来了 7808738
- 3 东部战区:开火!开火!全部命中! 7712893
- 4 2026年这些民生政策将惠及百姓 7616985
- 5 小学食堂米线过期2.5小时被罚5万 7519709
- 6 解放军喊话驱离台军 原声曝光 7428214
- 7 为博流量直播踩烈士陵墓?绝不姑息 7327605
- 8 每月最高800元!多地发放养老消费券 7238391
- 9 数字人民币升级 1月1日起将计付利息 7141831
- 10 2026年1月1日起 一批新规将施行 7040675








AI100
