Is is possible to mix "user define weights" with "torch.nn.Module"

1

Thanks for you attention. I am not quite familiar with PyTorch and wondering is it possible to mix weights defined by my self (not weights like that in torch.nn.conv2d, but like self.filter in the below example) into torch.nn.Module.

If that is possible. I have two questions:

  1. How can I bring it to the model.parameter and bring it to cuda?
  2. How can I add it to optimizer?

Thanks in advance! :grinning:

A sample,

class myModel(torch.nn.Module):
    def __init__(self):
        super(myModel, self).__init__()
        self.filter = torch.randn(8, 4, 3, 3)
    def forward(self, x):
        return torch.nn.functional.conv2d(x, self.filter)

inputs = torch.randn(1, 4, 5, 5).cuda()
model = myModel().cuda() # it doesn't work (1)
optimizer = torch.optim.SGD([model.filter], lr=0.01, momentum=0.9) # will it work? (2)

# result
>>> [i for i in model.parameters()]
[]
2

torch.nn.Parameter

1 Like
3

Coooool! That works.

here’s my testing code

class myModel(torch.nn.Module):
    def __init__(self):
        super(myModel, self).__init__()
        self.filter = torch.nn.Parameter(torch.randn(8, 4, 3, 3))

    def forward(self, x):
        return torch.nn.functional.conv2d(x, self.filter)

inputs = torch.randn(1, 4, 5, 5).cuda()
model = myModel().cuda()

gt = torch.randn(1, 8, 3, 3).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

for _ in range(10):
    optimizer.zero_grad()
    result = model(inputs)
    loss = torch.nn.functional.mse_loss(result, gt)
    print(loss)
    loss.backward(retain_graph=True)
    optimizer.step()