欢迎访问宙启技术站
智能推送

使用autograd构建神经网络模型的示例代码

发布时间:2024-01-11 07:36:57

使用autograd构建神经网络模型的示例代码:

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable

# 构建数据集
# 输入x为一个三维向量,输出y为一个二维向量
np.random.seed(0)
torch.manual_seed(0)

x = np.random.rand(100, 3)
w1 = np.array([[1, -2, 1],
               [-3, 2, 1]]).T
w2 = np.array([[1, 2]]).T
y = np.matmul(np.matmul(x, w1), w2)

# 转换为Tensor对象,并使用Variable包装
x = Variable(torch.Tensor(x), requires_grad=False)
y = Variable(torch.Tensor(y), requires_grad=False)


# 定义神经网络模型
class Model(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

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

# 创建模型对象
input_dim = 3
hidden_dim = 10
output_dim = 2
model = Model(input_dim, hidden_dim, output_dim)

# 定义优化器和损失函数
learning_rate = 0.01
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
criterion = nn.MSELoss()

# 训练模型
num_epochs = 100
for epoch in range(num_epochs):
    # 前向传播
    outputs = model(x)
    loss = criterion(outputs, y)

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

    # 打印训练信息
    if epoch % 10 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch, num_epochs, loss.item()))

# 测试模型
x_test = Variable(torch.Tensor([[0.1, 0.2, 0.3]]))
y_test = model(x_test)
print('Prediction for [{:.1f}, {:.1f}, {:.1f}]: {}'.format(x_test[0][0], x_test[0][1], x_test[0][2], y_test.data.numpy()))

上述代码使用autograd构建了一个简单的两层神经网络模型。首先,构建了一个输入维度为3、隐藏层维度为10、输出维度为2的模型对象。模型的前向传播使用ReLU作为激活函数,最后一层没有激活函数。优化器选择了随机梯度下降(SGD),损失函数选择了均方误差(MSE)。在训练过程中,对模型进行了100次迭代,每10次打印一次训练信息。最后,使用训练好的模型对一个测试样本进行预测。