pytorch自定义重量

时间:2019-11-20 04:45:58

标签: pytorch

我有一个网络

class Net(nn.Module)

和两个不同的权重w0w1(将所有图层的权重连接到一个向量中)。现在,我想在连接w0w1的线上优化网络,这意味着权重将采用theta * w0 + (1-theta) * w1的形式。因此,现在我要优化的参数不再是权重本身,而是theta

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

NetOnLine(nn.Module)

我应该如何编写forward(self, X)函数?

1 个答案:

答案 0 :(得分:0)

您可以在网络中将参数theta定义为nn.Parameter。您将以与普通方式相同的方式定义前向功能-通过所需的层或操作传递数据,然后将其返回。

这是一个最小的示例,我训练一个“网络”以学习将张量乘以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