PyTorch - Change weights of Conv2d

Viewed 499

For some reason, I cannot seem to assign all the weights of a Conv2d layer in PyTorch - I have to do it in two steps. Can anyone help me with what I am doing wrong?

layer = torch.nn.Conv2d(in_channels=1, out_channels=2, kernel_size=(2,2), stride=(2,2))
layer.state_dict()['weight']

gives me a tensor of size (2,1,2,2)

tensor([[[[ 0.4738, -0.2197],
      [-0.3436, -0.0754]]],

    [[[ 0.1662,  0.4098],
      [-0.4306, -0.4828]]]])

When I try to assign weights like so

layer.state_dict()['weight'] = torch.tensor([
  [[[ 1, 2],
    [3, 4]]],

  [[[-1, -2],
    [-3, -4]]]
])

the weights don't change. However, if I do something like this

layer.state_dict()['weight'][0] = torch.tensor([
  [[[1, 2],
    [3, 4]]],
])
layer.state_dict()['weight'][1] = torch.tensor([
  [[[-1, -2],
    [-3, -4]]],
])

The weights change. Why is this so?

1 Answers

I'm not sure about why you can't directly assign them but the more proper way to achieve what you're trying to do would be

layer.load_state_dict({'weight': torch.tensor([[[[0.4738, -0.2197],
                      [-0.3436, -0.0754]]],
                    [[[0.1662, 0.4098],
                      [-0.4306, -0.4828]]]])}, strict=False)
Related