当前位置：首页 > article >正文

PyTorch实战-模拟线性函数和非线性函数

article 2024/12/23 4:00:09

1 需求

需求：模拟y=x

需求：模拟y=x^2

2 接口

3.* y=x

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt


# 自定义数据集类
class MyDataset(Dataset):
    def __init__(self, start=-10, end=10, num_points=100):
        self.x = torch.linspace(start, end, num_points).view(-1, 1)
        self.y = self.x

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]


# 定义一个简单的线性模型，y = x 可以表示为 y = 1 * x + 0，这里使用线性层来模拟
class LinearModel(nn.Module):
    def __init__(self):
        super(LinearModel, self).__init__()
        # 输入维度为 1，输出维度为 1
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)


# 实例化模型
model = LinearModel()

# 定义损失函数，这里使用均方误差损失
criterion = nn.MSELoss()
# 定义优化器，使用随机梯度下降，学习率为 0.01
optimizer = optim.SGD(model.parameters(), lr=0.01)


# 实例化自定义数据集
dataset = MyDataset()
# 实例化 DataLoader，batch_size 为 10
dataloader = DataLoader(dataset, batch_size=10, shuffle=True)


# 开始训练
num_epochs = 1000
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (inputs, targets) in enumerate(dataloader):
        # 前向传播
        outputs = model(inputs)
        # 计算损失
        loss = criterion(outputs, targets)
        # 清空梯度
        optimizer.zero_grad()
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()
        running_loss += loss.item()
    if (epoch + 1) % 100 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss / len(dataloader):.4f}')


# 测试模型
# 生成测试数据，这里使用和训练数据相同的范围，也可以使用新的范围
x_test = torch.linspace(-10, 10, 100).view(-1, 1)
y_test = model(x_test)


# 可视化结果
plt.plot(dataset.x.detach().numpy(), dataset.y.detach().numpy(), label='True')
plt.plot(x_test.detach().numpy(), y_test.detach().numpy(), label='Predicted')
plt.legend()
plt.xlabel('x')
plt.ylabel('y')
plt.title('y = x')
plt.show()

3.* y=x^2

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt


# 自定义数据集类
class MyDataset(Dataset):
    def __init__(self, start=-10, end=10, num_points=100):
        self.x = torch.linspace(start, end, num_points).view(-1, 1)
        # 计算 y = x^2
        self.y = self.x ** 2

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]


# 定义一个简单的神经网络模型
class QuadraticModel(nn.Module):
    def __init__(self):
        super(QuadraticModel, self).__init__()
        # 一个简单的神经网络，包含一个隐藏层
        self.fc1 = nn.Linear(1, 10)
        self.fc2 = nn.Linear(10, 1)

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


# 实例化模型
model = QuadraticModel()

# 定义损失函数，这里使用均方误差损失
criterion = nn.MSELoss()
# 定义优化器，使用随机梯度下降，学习率为 0.01
optimizer = optim.SGD(model.parameters(), lr=0.01)


# 实例化自定义数据集
dataset = MyDataset()
# 实例化 DataLoader，batch_size 为 10
dataloader = DataLoader(dataset, batch_size=10, shuffle=True)


# 开始训练
num_epochs = 1000
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (inputs, targets) in enumerate(dataloader):
        # 前向传播
        outputs = model(inputs)
        # 计算损失
        loss = criterion(outputs, targets)
        # 清空梯度
        optimizer.zero_grad()
        # 反向传播
        loss.backward()
        # 更新参数
        optimizer.step()
        running_loss += loss.item()
    if (epoch + 1) % 100 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss / len(dataloader):.4f}')


# 测试模型
# 生成测试数据，这里使用和训练数据相同的范围，也可以使用新的范围
x_test = torch.linspace(-10, 10, 100).view(-1, 1)
y_test = model(x_test)


# 可视化结果
plt.plot(dataset.x.detach().numpy(), dataset.y.detach().numpy(), label='True')
plt.plot(x_test.detach().numpy(), y_test.detach().numpy(), label='Predicted')
plt.legend()
plt.xlabel('x')
plt.ylabel('y')
plt.title('y = x^2')
plt.show()