Error when running in test mode in PyTorch-Lightning Trainer

[angadkalra]

Trying to use medcam with 3D ResNet-101 model implemented in PyTorch-Lightning. Data is CT scans. Does medcam work when using "ddp" or "dp" backend? Currently getting this error:

RuntimeError: invalid gradient at index 0 - expected type TensorOptions(dtype=c10::Half, device=cuda:0, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt)) but got TensorOptions(dtype=c10::Half, device=cpu, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))

This is my code to run inference:

# Load model from ckpt
    model = load_model(args.ckpt_path, args.model_depth, input_type=args.input_shape)
    if args.explain:
        print("\nGenerating attenion maps...")
        model = medcam.inject(model, output_dir=os.path.join(CONFIG.LTRC_PATH, "attention_maps"), save_maps=True)

# Infer on test set
    trainer = pl.Trainer(accelerator="ddp", gpus=args.num_gpus, precision=16, logger=None)
    trainer.test(model=model, datamodule=datamodule, verbose=True)

Using 1 GPU (V100 with 16gb mem) and running on google cloud compute engine VM.

Thanks!

[Karol-G]

Hi,

I never tested medcam with DataParallel or DistributedDataParallel so I don't know. Does it work without any of them? Plus, can you post the full stack trace so I can see where the exception is raised?

Best Karol

[angadkalra]

Here is full stack trace:

File "pytorch_infer.py", line 57, in <module>
    main(args)
  File "pytorch_infer.py", line 38, in main
    trainer.test(model=model, datamodule=datamodule, verbose=True)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 708, in test
    results = self.__test_given_model(model, test_dataloaders)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 773, in __test_given_model
    results = self.fit(model)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 440, in fit
    results = self.accelerator_backend.train()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/dp_accelerator.py", line 97, in train
    results = self.train_or_test()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py", line 64, in train_or_test
    results = self.trainer.run_test()
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 617, in run_test
    eval_loop_results, _ = self.run_evaluation(test_mode=True)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 568, in run_evaluation
    output = self.evaluation_loop.evaluation_step(test_mode, batch, batch_idx, dataloader_idx)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 169, in evaluation_step
    output = self.trainer.accelerator_backend.test_step(args)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/dp_accelerator.py", line 118, in test_step
    output = self.training_step(args)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/accelerators/dp_accelerator.py", line 108, in training_step
    output = self.trainer.model(*args)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/cuda/amp/autocast_mode.py", line 135, in decorate_autocast
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 82, in forward
    return self.module.test_step(*inputs[0], **kwargs[0])
  File "/imvaria/fibresolve_base/fibresolve/models/torch_resnet_base.py", line 53, in test_step
    out = self._get_loss_preds(batch)
  File "/imvaria/fibresolve_base/fibresolve/models/torch_resnet_base.py", line 31, in _get_loss_preds
    logits = torch.squeeze(self(X), dim=1)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/medcam/medcam_inject.py", line 199, in forward
    self.test_run(batch, internal=True)
  File "/opt/conda/lib/python3.7/site-packages/medcam/medcam_inject.py", line 229, in test_run
    _ = self.medcam_dict['model_backend'].generate_attention_map(batch, None)
  File "/opt/conda/lib/python3.7/site-packages/medcam/backends/base.py", line 22, in generate_attention_map
    self.backward(label=label)
  File "/opt/conda/lib/python3.7/site-packages/medcam/backends/base.py", line 42, in backward
    self.logits.backward(gradient=self.mask, retain_graph=self.retain_graph)
  File "/opt/conda/lib/python3.7/site-packages/torch/tensor.py", line 185, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/opt/conda/lib/python3.7/site-packages/torch/autograd/__init__.py", line 127, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: invalid gradient at index 0 - expected type TensorOptions(dtype=c10::Half, device=cuda:0, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt)) but got TensorOptions(dtype=c10::Half, device=cpu, layout=Strided, requires_grad=false (default), pinned_memory=false (default), memory_format=(nullopt))

Could you please look into adding functionality for DP and DDP?

Thanks!

[Karol-G]

Thanks for the stack trace. Can you test if you get the same error if you turn of DP/DDP and use normal precision?

I would also recommend to use medcam after you trained your model not before. So do something like this:

# Load model from ckpt
model = load_model(args.ckpt_path, args.model_depth, input_type=args.input_shape)

# Infer on test set
trainer = pl.Trainer(accelerator="ddp", gpus=args.num_gpus, precision=16, logger=None)

if args.explain:
    print("\nGenerating attenion maps...")
    model = medcam.inject(model, output_dir=os.path.join(CONFIG.LTRC_PATH, "attention_maps"), save_maps=True)

trainer.test(model=model, datamodule=datamodule, verbose=True)

Best Karol