Pytorch 自定义权重

Question

我有一个网络

class Net(nn.Module)

还有两个不同的权重w0 和w1（将所有层的权重连接成一个向量）。现在我想优化连接w0 和w1 的线上的网络，这意味着权重将具有theta * w0 + (1-theta) * w1 的形式。所以现在我要优化的参数不再是权重本身，而是theta。

我该如何实现呢？在 Pytorch 中，如何将参数定义为theta，并将权重设置为我想要的形式。具体来说，如果我创建一个新类

NetOnLine(nn.Module)

forward(self, X)函数应该怎么写？

Answer 1

您可以在您的网络中将参数theta 定义为nn.Parameter。您可以像往常一样定义 forward 函数 - 通过所需的层或操作传递数据，然后返回它。

这是一个最小的例子，我训练一个“网络”来学习将张量乘以 2：

import numpy as np
import torch


class SampleNet(torch.nn.Module):
    def __init__(self):
        super(SampleNet, self).__init__()

        self.theta = torch.nn.Parameter(torch.rand(1))

    def forward(self, x):
        x = x * self.theta.expand_as(x)  # expand_as() to match sizes
        return x


train_data = np.random.rand(1000, 10)
train_data[:, 5:] = 2 * train_data[:, :5]
train_data = torch.Tensor(train_data)

sample_net = SampleNet()

optimizer = torch.optim.Adam(params=sample_net.parameters())
mse_loss = torch.nn.MSELoss()

for epoch in range(5):
    for data in train_data:
        x = data[:5]
        y = data[5:]

        optimizer.zero_grad()
        prediction = sample_net(x)
        loss = mse_loss(y, prediction)
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch}, Loss {loss.data.item()}")
print(f"Learned theta: {sample_net.theta.data.item()}")

打印出来

Epoch 0, Loss 0.03369491919875145
Epoch 1, Loss 0.0018534092232584953
Epoch 2, Loss 1.2343853995844256e-05
Epoch 3, Loss 2.2044337466553543e-09
Epoch 4, Loss 4.0527581290916714e-12
Learned theta: 1.999994158744812