为什么火把闪电configure_optimizer扔AssertionError: param组一定是一个小块？

Question

我已经建立了多个火把闪电项目在过去和建立一个新的快速演示项目，我偶然发现这个奇怪的错误，不知怎么我无法摆脱它。

这是我的模型文件的相关部分。

class TSModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(3, 10, kernel_size=(3, 3), padding=(1, 1)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))
        )

        self.classifier = nn.Sequential(
            nn.Linear(10*16*16, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )

        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        N = x.shape[0]
        x = self.backbone(x)
        x = x.view(N, -1)
        return self.classifier(x)

    def configure_optimizers(self):
        params = [p for p in self.parameters() if p.requires_grad]
        return torch.optim.AdamW(self.parameters())

但是，在启动培训过程时，程序将退出，并引发以下内容：

Traceback (most recent call last):
  File "/torchserve-example/main.py", line 25, in <module>
    ts_train()
  File "/torchserve-example/main.py", line 21, in ts_train
    trainer.fit(model, datamodule)
  File ".local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 458, in fit
    self._run(model)
  File ".local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 715, in _run
    self.accelerator.setup(self, model)  # note: this sets up self.lightning_module
  File ".local/lib/python3.8/site-packages/pytorch_lightning/accelerators/cpu.py", line 39, in setup
    return super().setup(trainer, model)
  File ".local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 92, in setup
    self.setup_optimizers(trainer)
  File ".local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 374, in setup_optimizers
    optimizers, lr_schedulers, optimizer_frequencies = self.training_type_plugin.init_optimizers(
  File ".local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 190, in init_optimizers
    return trainer.init_optimizers(model)
  File ".local/lib/python3.8/site-packages/pytorch_lightning/trainer/optimizers.py", line 34, in init_optimizers
    optim_conf = model.configure_optimizers()
  File "/torchserve-example/model.py", line 52, in configure_optimizers
    return torch.optim.AdamW(self.parameters())
  File ".local/lib/python3.8/site-packages/torch/optim/adamw.py", line 47, in __init__
    super(AdamW, self).__init__(params, defaults)
  File ".local/lib/python3.8/site-packages/torch/optim/optimizer.py", line 55, in __init__
    self.add_param_group(param_group)
  File ".local/lib/python3.8/site-packages/torch/optim/optimizer.py", line 242, in add_param_group
    assert isinstance(param_group, dict), "param group must be a dict"
AssertionError: param group must be a dict

当我在print(type(params[0]))中执行configure_optimizers时，它会将<class 'torch.nn.parameter.Parameter'>打印到stdout。知道这里出了什么问题吗？

注意:由于这个错误发生在优化器的初始化过程中，这很可能与火把闪电没有直接关系，这也是为什么我也将py手电作为一个标签。

Answer 1

在图书馆代码中，我发现：

# if not isinstance(param_groups[0], dict):
#             param_groups = [{'params': param_groups}]

在注释这一点时，一切正常工作。我没有回答这个问题，因为更改基础库或将此代码部分复制到我的文件中并不是一个很好的解决方案。

Answer 2

实际上，这一行代码是错误的：

def configure_optimizers(self):
        params = [p for p in self.parameters() if p.requires_grad]
        return torch.optim.AdamW(self.parameters())

您传递的是params，而不是self.parameters()，因为这样做很好。

使用像这样创建的params，本质上是传递带有生成器的列表，不是dict的实例。

在PyTorch中，可以通过列表中包含的dicts传递具有不同学习速率等的多个不同参数。在PyTorch的API中，这就是您的params看起来的样子。