如何遍历 Pytorch 中未知形状的整个张量？

Question

假设我有一个任意的 NN.module。 我想遍历这个 neural.network 的所有可学习权重，然后修改每个权重。

例如我想做这样的事情（其中 model 是 Re.net50）：

with torch.no_grad():
  for param in model.parameters():
    print(type(param), param.size())
    param[0][0][0][0] = 0.5
    print(param[0][0][0][0])
    break

<class 'torch.nn.parameter.Parameter'> torch.Size([64, 3, 7, 7])
tensor(0.5000, device='cuda:0', requires_grad=True)

但是，我想在不知道参数尺寸的情况下对所有权重执行此操作。 无论如何要这样做还是我需要做一些递归？

Answer 1

您可以通过展平来迭代任意大小和维数的张量。

for param in model.parameters():
  shape = param.shape # store original shape
  flat = param.flatten()
  for i in range(len(flat)):
      flat[i] = ... modify item, unclear what you want to do...

  # return to original shape
  param = flat.view(shape)