pytorch中Parameter函数用法示例

PyTorch中Parameter函数用法示例

在PyTorch中，Parameter函数是一个特殊的张量，它被自动注册为模型的可训练参数。本文将介绍Parameter函数的用法，并演示两个示例。

示例一：使用Parameter函数定义可训练参数

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.weight = nn.Parameter(torch.randn(10, 5))

    def forward(self, x):
        return torch.matmul(x, self.weight)

在上述代码中，我们首先定义了一个MyModel类，继承自nn.Module。在__init__()方法中，我们使用nn.Parameter函数定义了一个10x5的可训练参数weight。在forward()方法中，我们将输入x与weight进行矩阵乘法，并返回输出。

示例二：使用Parameter函数更新可训练参数

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.weight = nn.Parameter(torch.randn(10, 5))

    def forward(self, x):
        return torch.matmul(x, self.weight)

# 实例化模型
model = MyModel()

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

# 训练模型
for epoch in range(10):
    inputs = torch.randn(3, 10)
    labels = torch.randn(3, 5)
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, 10, loss.item()))

在上述代码中，我们首先实例化了MyModel类，并定义了损失函数和优化器。然后，我们使用for循环训练模型，并使用optimizer.step()函数更新可训练参数。需要注意的是，我们使用model.parameters()函数获取模型的可训练参数。

结论

总之，在PyTorch中，我们可以使用nn.Parameter函数定义可训练参数，并使用optimizer.step()函数更新可训练参数。需要注意的是，不同的模型可能会有不同的可训练参数，因此需要根据实际情况进行调整。