温馨提示×

在PyTorch中如何进行模型权重的正则化

小樊
107
2024-03-05 19:47:07
栏目: 编程语言

在PyTorch中,可以使用torch.nn.Module类中的parameters()方法来获取模型的权重参数,然后使用正则化方法来对这些参数进行约束。下面是一个示例代码,演示了如何对模型的权重进行L2正则化:

import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 1)

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

# 创建模型实例
model = Net()

# 定义L2正则化参数
l2_lambda = 0.01

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

# 训练模型
for epoch in range(100):
    optimizer.zero_grad()
    
    # 正向传播
    output = model(torch.randn(1, 10))
    loss = criterion(output, torch.randn(1, 1))
    
    # 添加L2正则化项
    l2_reg = torch.tensor(0.)
    for param in model.parameters():
        l2_reg += torch.norm(param)
    
    loss += l2_lambda * l2_reg
    
    # 反向传播
    loss.backward()
    optimizer.step()

在上述示例中,我们首先定义了一个简单的神经网络模型Net,然后创建了模型实例。在训练循环中,我们使用optimizer.zero_grad()来清除之前的梯度,然后进行正向传播和计算损失。接着,我们计算所有权重参数的L2范数,并将其添加到损失函数中作为正则化项。最后,我们进行反向传播和更新模型参数。

0