KeyError: Caption

[dreamhhy]

Does anyone have such issue? I use ImageNet2012 to train diffusion model.

Traceback (most recent call last): File "main.py", line 726, in trainer.fit(model, data) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 553, in fit self._run(model) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 918, in _run self._dispatch() File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 986, in _dispatch self.accelerator.start_training(self) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 92, in start_training self.training_type_plugin.start_training(trainer) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 161, in start_training self._results = trainer.run_stage() File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 996, in run_stage return self._run_train() File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1031, in _run_train self._run_sanity_check(self.lightning_module) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1115, in _run_sanity_check self._evaluation_loop.run() File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 111, in run self.advance(*args, **kwargs) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 110, in advance dl_outputs = self.epoch_loop.run( File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 111, in run self.advance(*args, **kwargs) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 110, in advance output = self.evaluation_step(batch, batch_idx, dataloader_idx) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 154, in evaluation_step output = self.trainer.accelerator.validation_step(step_kwargs) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 211, in validation_step return self.training_type_plugin.validation_step(*step_kwargs.values()) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 386, in validation_step return self.model(*args, **kwargs) File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl result = self.forward(*input, **kwargs) File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 619, in forward output = self.module(*inputs[0], **kwargs[0]) File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl result = self.forward(*input, **kwargs) File "/root/miniconda3/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 93, in forward output = self.module.validation_step(*inputs, **kwargs) File "/root/miniconda3/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context return func(*args, **kwargs) File "/root/autodl-tmp/latent-diffusion-github/ldm/models/diffusion/ddpm.py", line 359, in validation_step _, loss_dict_no_ema = self.shared_step(batch) File "/root/autodl-tmp/latent-diffusion-github/ldm/models/diffusion/ddpm.py", line 868, in shared_step x, c = self.get_input(batch, self.first_stage_key) File "/root/miniconda3/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context return func(*args, **kwargs) File "/root/autodl-tmp/latent-diffusion-github/ldm/models/diffusion/ddpm.py", line 670, in get_input xc = batch[cond_key] KeyError: 'caption'

[GrandpaXun242]

Your Dataset class is wrong. It should contain the key caption in batch dict.