Pytorch customize weight
I have a network

class Net(nn.Module)


and two different weights w0 and w1 (concatenate weights of all layers into a vector). Now I want to optimize the network on the line connecting w0 and w1, which means that the weight will have the form theta  * w0 + (1-theta) * w1. So now the parameter I want to optimize is no longer the weight itself, but the theta. 

How can I implement this? In Pytorch, how can I define the parameter to be theta, and set the weight to be form I want. To be specific, if I create a new class 

NetOnLine(nn.Module)


how should I write the forward(self, X) function?

You can define the parameter theta in your net as an nn.Parameter. You'd define the forward function the same way as normal - pass the data through the layers or operations you want and then return it.

Here's a minimal example, where I train a "network" to learn to multiply a Tensor by 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()}")


which prints out

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

