一个因为迭代器产生的错误

python的生成器默认只会遍历一次，遍历一遍后，再进行取元素就不会返回任何值。

pytorch 的模型的 .parameters() 返回的是一个生成器，因此，在每次遍历时候都需要重新创建一次，或者使用别的方式遍历所有参数。

下面代码使用注释的一行和不注释的一行得到的结果不同。

使用 for param in self.paras: 这一行时候，输出结果：

grad_sum is:  tensor(0.)
grad_shape is:  torch.Size([485])

使用 for param in self.mlp.parameters(): 时，输出结果：

grad_sum is:  tensor(14.3077)
grad_shape is:  torch.Size([485])

import torch
import torch.nn as nn


class test_model(nn.Module):
    def __init__(self):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(10, 20),  # 10 * 20 + 20 = 220
            nn.ReLU(),
            nn.Linear(20, 10),  # 20 * 10 + 10 = 210
            nn.ReLU(),
            nn.Linear(10, 5),  # 10 * 5 + 5 = 55
        )
        self.paras = self.mlp.parameters()

        self.grad_index = []
        
        for param in self.paras:
        # for param in self.mlp.parameters():
        
            self.grad_index.append(param.data.numel())
        self.grad_dim = sum(self.grad_index)

    def get_grad_vec(self):
        count = 0
        begin_idx = 0
        grad = torch.zeros(self.grad_dim)
        
        for param in self.paras:
        # for param in self.mlp.parameters():
        
            para_num = self.grad_index[count]
            if param.grad is not None:
                grad[begin_idx : begin_idx + para_num] = param.grad.data.view(-1)
            begin_idx += para_num
            count += 1

        return grad

    def forward(self, x):
        return self.mlp(x)


xxx = torch.rand(10)
model = test_model()
model.zero_grad()
output = model(xxx)
output.sum().backward()

print("grad_sum is: ", model.get_grad_vec().sum())
print("grad_shape is: ", model.get_grad_vec().shape)