python pytorch 实战篇：a+b数字模型-优快云博客

本文链接：https://blog.youkuaiyun.com/RedStone114514/article/details/148642094

前言

看之前必须得去看我的理论篇，不然跟不上：python pytorch 从入门到精通通俗易懂
实战篇终于出来啦！！！数字版a+b模型！
全部整体代码结尾自取

构思

输入两个数字 a , b ,输出他们的和

定义模型

我们就是输入两个数，输出一个数，那么我们就可以定义成：
2 (输入层) -> 4 (隐藏层) -> 1 (输出层)
所以我们可以得出

# define.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class Adder(nn.Module):
    def __init__(self):
        super(Adder,self).__init__()
        
        self.layer = nn.Sequential(
            nn.Linear(2, 4),
            nn.ReLU(),
            nn.Linear(4, 1)
        )

    def forward(self, x):
        # 确保输入是二维张量
        x = x.unsqueeze(0).float()
        out = self.layer(x)
        # 将输出张量降维为一维张量
        return out.squeeze(0)

x.unsqueeze(0)：unsqueeze(0) 方法的作用是在输入张量 x 的第 0 个维度插入一个新的维度。若输入 x 是一维张量 (input_features)，那么 x.unsqueeze(0) 之后就会变成二维张量 (1, input_features)，这里的 1 代表批量大小为 1。
out.squeeze(0)：squeeze(0) 方法与 unsqueeze(0) 相反，它会移除张量中维度大小为 1 的第 0 个维度。由于网络输出是二维张量 (1, output_features)，经过 squeeze(0) 后就会变回一维张量 (output_features)。

定义数据集

我们可以把数据信息定义成这样：
((a,b,add [a,b的和]), ...)
我们就可以轻松地写出：

# define.py

class Dataset:
    def __init__(self, data):
        self.inp = torch.tensor(
            [[a, b] for a, b, _ in data],
            dtype=torch.int
        )

        self.out = torch.tensor(
            [[add] for _, _, add in data],
            dtype=torch.int
        )

    def __len__(self):
        return len(self.inp)
    
    def __getitem__(self, index):
        return self.inp[index], self.out[index]

__init__：定义输入输出列表
__len__：获取长度
__getitem__：获取搜引，这得要看你后面怎么写来定义

训练

我们的任务归属于回归任务所以我们使用
损失器：MSELoss
优化器：Adam

# train.py

from define import *

import torch
import torch.nn as nn
import torch.functional as F
import torch.optim as optim

import os

import random

# 超参数
EPOCHS = 1000
LR = 0.001

model = Adder() # 模型

# 转移设备
torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 损失函数
loss_fn = nn.MSELoss()

# 优化器
optimizer = optim.Adam(model.parameters(), lr=LR)

# 数据
data = [(1,2,3),(4,5,9),(6,7,13)]
for _ in range(1000):
    a = random.randint(0, 100)
    b = random.randint(0, 100)
    add = a + b
    data.append((a, b, add))
dataset = Dataset(data)

# 训练
for epochin range(EPOCHS):
    for x, y in dataset:
        # 前向传播
        y_pred = model(x)

        # 计算损失
        loss = loss_fn(y_pred, y.float())

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

	# 每 10 步 输出误差
    if epoch % 10 == 0:
        print(loss.item())

# 保存
torch.save(model, 'adder.pth')

优化（选看）

我们假如像要在原有的继续训练，我们可以加上一个判断：

# 模型
if os.path.exists('adder.pth'):
    model = torch.load('adder.pth', weights_only=False)
else:
    model = Adder()

训练时Ctrl+C退出自动保存

try:
    for epoch in range(EPOCHS):
    
    ... # 训练过程
    
except KeyboardInterrupt:
    pass
torch.save(model, 'adder.pth')
print("save")

运行

我们输入两个数字，然后输出他们的推理结果

from define import *

import torch

# 创建模型实例
model = Adder()
model = torch.load('adder.pth', weights_only=False)

while True:
    with torch.no_grad():  # 关闭梯度计算加速推理
        a = input("a:")
        b = input("b:")

        # 确保输入为浮点数
        input_tensor = torch.tensor([[int(a), int(b)]], dtype=torch.int)    # 输入必须是张量
        prediction = model(input_tensor)             # 输出也是张量
        print(round(prediction.item()))    # 输出

round四舍五入：因为输出的数字会有一点点的偏差，例如 0.00几的，我们可以使用四舍五入来取整数

完成

好！现在完成了！！！
其实人工智能并不是那么复杂，只是一个模拟人脑的东西，只要训练数据、批次够多，什么都可以做

最终代码（含优化）

define.py

这是定义代码

import torch
import torch.nn as nn
import torch.nn.functional as F

class Adder(nn.Module):
    def __init__(self):
        super(Adder,self).__init__()
        
        self.layer = nn.Sequential(
            nn.Linear(2, 4),
            nn.ReLU(),
            nn.Linear(4, 1)
        )

    def forward(self, x):
        # 确保输入是二维张量
        x = x.unsqueeze(0).float()
        out = self.layer(x)
        # 将输出张量降维为一维张量
        return out.squeeze(0)

# 数据: ((a,b,a+b), ...)

class Dataset:
    def __init__(self, data):
        self.inp = torch.tensor(
            [[a, b] for a, b, _ in data],
            dtype=torch.int
        )

        self.out = torch.tensor(
            [[add] for _, _, add in data],
            dtype=torch.int
        )

    def __len__(self):
        return len(self.inp)
    
    def __getitem__(self, index):
        return self.inp[index], self.out[index]

train.py

训练代码

from define import *

import torch
import torch.nn as nn
import torch.functional as F
import torch.optim as optim

import os

import random

# 超参数
EPOCHS = 1000
LR = 0.001

# 模型
if os.path.exists('adder.pth'):
    model = torch.load('adder.pth', weights_only=False)
else:
    model = Adder()

torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("run on",("cuda" if torch.cuda.is_available() else "cpu"))

# 损失函数
loss_fn = nn.MSELoss()

# 优化器
optimizer = optim.Adam(model.parameters(), lr=LR)

# 数据
data = [(1,2,3),(4,5,9),(6,7,13)]
for _ in range(1000):
    a = random.randint(0, 100)
    b = random.randint(0, 100)
    add = a + b
    data.append((a, b, add))

dataset = Dataset(data)

# 训练
try:
    for epoch in range(EPOCHS):
        for x, y in dataset:
            # 前向传播，只传入 x
            y_pred = model(x)

            # 计算损失
            loss = loss_fn(y_pred, y.float())

            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # if (loss.item()==0):
        #     print("end!! 0.0!!!")
        #     break
        # 打印损失
        if epoch % 10 == 0:
            print(loss.item())
except KeyboardInterrupt:
    pass

torch.save(model, 'adder.pth')
print("saving")

run.py

运行模型代码

from define import *

import torch
# import torch.nn as nn
# import torch.functional as F

# 创建模型实例
model = Adder()
model = torch.load('adder.pth', weights_only=False)

while True:
    with torch.no_grad():  # 关闭梯度计算加速推理
        a = input("a:")
        b = input("b:")

        # 确保输入为浮点数
        input_tensor = torch.tensor([[int(a), int(b)]], dtype=torch.int)    # 输入必须是张量
        prediction = model(input_tensor)             # 输出也是张量
        print(round(prediction.item()))    # 输出