[BUG] aten::max_pool3d_with_indices not implemented #326
Comments
|
Hi Conrad! Thanks a lot for giving it a go. It never clicked in my head that it was the max pool 3D, and we probably tested it using a 2D dataset... We will need to figure out a work around... In the error it states "As a temporary fix, you can set the environment variable Have you tried this? |
|
Hi Joran, I haven't tried it, but because it kind of defeats the purpose... Obviously it is possible to run everything on CPU; what I'm interested in though is the ballpark of the time it would take to process different sized images on macOS. I took a quick look at the PyTorch issue for maxpool_3d_with_indices. Maybe I muster the motivation to give it a go over the holidays... |
|
My understanding of the flag is that it only runs that single operation on CPU (or any operation that is still not supported), not the whole training. I am not familiar with the It is on our todo list to test it out quantitatively, but the end of the year is a bit tough... |
|
Hi @conradkun, I have just taken the time to do a quick comparison of running 3D N2V on mac with only CPU vs MPS with To do this I ran the following script (use from pathlib import Path
import timeit
import tifffile
from careamics import CAREamist
from careamics.config import create_n2v_configuration
from careamics_portfolio import PortfolioManager
# instantiate data portfolio manage
portfolio = PortfolioManager()
# and download the data
root_path = Path("./data")
file = portfolio.denoising.Flywing.download(root_path)
train_image = tifffile.imread(file[0])
n_epochs = 1
config = create_n2v_configuration(
experiment_name="flywing_n2v",
data_type="array",
axes="ZYX",
patch_size=(16, 64, 64),
batch_size=2,
num_epochs=n_epochs,
augmentations=[], # remove augmentations
)
# train
n_repeats = 10
total_train_time = timeit.timeit(
lambda: CAREamist(source=config).train(
train_source=train_image,
val_percentage=0.0,
val_minimum_split=10, # use 10 patches as validation
),
number=n_repeats,
)
average_train_time = total_train_time/n_repeats
print(f"Average time taken to train {n_epochs} epoch(s): {average_train_time :.2f}s")For CPU only the print out was For MPS with This is still about a 4-fold speed-up so the overhead of transferring the tensors between the GPU and CPU doesn't seem too bad. Hopefully the MPS implementation of Let me know if you have any questions 😊 Specs |
|
And if anyone runs into that issue as well, and stumble upon that message, please thumb up the following message: pytorch/pytorch#77764 (comment) In the off-chance that it will push the operation up the priority list! |
Describe the bug
Unfortunately, I was able to reproduce the
mpsbug whereaten::max_pool3d_with_indicesis missing on macOS GPU acceleration. Indeed, this operation is shown as "In Process" in the PyTorch mps Ops tracker, meaning it is not implemented yet. This error seems to (understandably) only pop up when trying to process 3D images usingmpson macOS.To Reproduce
I created a new environment from scratch, following the guidelines in the CAREamics installation guide.
Then I just tried to train on a 3D image, albeit one with a custom file extension (and hence custom
read_npyfunction). I'll be happy to try with a.tifffile if you point me to one.Code snippet allowing reproducing the behaviour:
Error message:
{ "name": "NotImplementedError", "message": "The operator 'aten::max_pool3d_with_indices' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.", "stack": "--------------------------------------------------------------------------- NotImplementedError Traceback (most recent call last) Cell In[32], line 18 16 engine = CAREamist(config) 17 start_time = time() ---> 18 engine.train(datamodule=data_module) 19 training_elapsed = time() - start_time File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/careamics/careamist.py:333, in CAREamist.train(self, datamodule, train_source, val_source, train_target, val_target, use_in_memory, val_percentage, val_minimum_split) 331 # train 332 if datamodule is not None: --> 333 self._train_on_datamodule(datamodule=datamodule) 335 else: 336 # raise error if target is provided to N2V 337 if self.cfg.algorithm_config.algorithm == SupportedAlgorithm.N2V.value: File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/careamics/careamist.py:410, in CAREamist._train_on_datamodule(self, datamodule) 407 self.trainer.limit_val_batches = 1.0 # 100% 409 # train --> 410 self.trainer.fit(self.model, datamodule=datamodule) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:538, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path) 536 self.state.status = TrainerStatus.RUNNING 537 self.training = True --> 538 call._call_and_handle_interrupt( 539 self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path 540 ) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:47, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs) 45 if trainer.strategy.launcher is not None: 46 return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs) ---> 47 return trainer_fn(*args, **kwargs) 49 except _TunerExitException: 50 _call_teardown_hook(trainer) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:574, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path) 567 assert self.state.fn is not None 568 ckpt_path = self._checkpoint_connector._select_ckpt_path( 569 self.state.fn, 570 ckpt_path, 571 model_provided=True, 572 model_connected=self.lightning_module is not None, 573 ) --> 574 self._run(model, ckpt_path=ckpt_path) 576 assert self.state.stopped 577 self.training = False File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:981, in Trainer._run(self, model, ckpt_path) 976 self._signal_connector.register_signal_handlers() 978 # ---------------------------- 979 # RUN THE TRAINER 980 # ---------------------------- --> 981 results = self._run_stage() 983 # ---------------------------- 984 # POST-Training CLEAN UP 985 # ---------------------------- 986 log.debug(f\"{self.__class__.__name__}: trainer tearing down\") File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1023, in Trainer._run_stage(self) 1021 if self.training: 1022 with isolate_rng(): -> 1023 self._run_sanity_check() 1024 with torch.autograd.set_detect_anomaly(self._detect_anomaly): 1025 self.fit_loop.run() File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1052, in Trainer._run_sanity_check(self) 1049 call._call_callback_hooks(self, \"on_sanity_check_start\") 1051 # run eval step -> 1052 val_loop.run() 1054 call._call_callback_hooks(self, \"on_sanity_check_end\") 1056 # reset logger connector File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/loops/utilities.py:178, in _no_grad_context.<locals>._decorator(self, *args, **kwargs) 176 context_manager = torch.no_grad 177 with context_manager(): --> 178 return loop_run(self, *args, **kwargs) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py:135, in _EvaluationLoop.run(self) 133 self.batch_progress.is_last_batch = data_fetcher.done 134 # run step hooks --> 135 self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter) 136 except StopIteration: 137 # this needs to wrap the `*_step` call too (not just `next`) for `dataloader_iter` support 138 break File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py:396, in _EvaluationLoop._evaluation_step(self, batch, batch_idx, dataloader_idx, dataloader_iter) 390 hook_name = \"test_step\" if trainer.testing else \"validation_step\" 391 step_args = ( 392 self._build_step_args_from_hook_kwargs(hook_kwargs, hook_name) 393 if not using_dataloader_iter 394 else (dataloader_iter,) 395 ) --> 396 output = call._call_strategy_hook(trainer, hook_name, *step_args) 398 self.batch_progress.increment_processed() 400 if using_dataloader_iter: 401 # update the hook kwargs now that the step method might have consumed the iterator File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:319, in _call_strategy_hook(trainer, hook_name, *args, **kwargs) 316 return None 318 with trainer.profiler.profile(f\"[Strategy]{trainer.strategy.__class__.__name__}.{hook_name}\"): --> 319 output = fn(*args, **kwargs) 321 # restore current_fx when nested context 322 pl_module._current_fx_name = prev_fx_name File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py:411, in Strategy.validation_step(self, *args, **kwargs) 409 if self.model != self.lightning_module: 410 return self._forward_redirection(self.model, self.lightning_module, \"validation_step\", *args, **kwargs) --> 411 return self.lightning_module.validation_step(*args, **kwargs) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/careamics/lightning/lightning_module.py:139, in FCNModule.validation_step(self, batch, batch_idx) 129 \"\"\"Validation step. 130 131 Parameters (...) 136 Batch index. 137 \"\"\" 138 x, *aux = batch --> 139 out = self.model(x) 140 val_loss = self.loss_func(out, *aux) 142 # log validation loss File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs) 1734 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc] 1735 else: -> 1736 return self._call_impl(*args, **kwargs) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module._call_impl(self, *args, **kwargs) 1742 # If we don't have any hooks, we want to skip the rest of the logic in 1743 # this function, and just call forward. 1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1745 or _global_backward_pre_hooks or _global_backward_hooks 1746 or _global_forward_hooks or _global_forward_pre_hooks): -> 1747 return forward_call(*args, **kwargs) 1749 result = None 1750 called_always_called_hooks = set() File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/careamics/models/unet.py:439, in UNet.forward(self, x) 425 def forward(self, x: torch.Tensor) -> torch.Tensor: 426 \"\"\" 427 Forward pass. 428 (...) 437 Output of the model. 438 \"\"\" --> 439 encoder_features = self.encoder(x) 440 x = self.decoder(*encoder_features) 441 x = self.final_conv(x) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs) 1734 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc] 1735 else: -> 1736 return self._call_impl(*args, **kwargs) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module._call_impl(self, *args, **kwargs) 1742 # If we don't have any hooks, we want to skip the rest of the logic in 1743 # this function, and just call forward. 1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1745 or _global_backward_pre_hooks or _global_backward_hooks 1746 or _global_forward_hooks or _global_forward_pre_hooks): -> 1747 return forward_call(*args, **kwargs) 1749 result = None 1750 called_always_called_hooks = set() File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/careamics/models/unet.py:124, in UnetEncoder.forward(self, x) 122 encoder_features = [] 123 for module in self.encoder_blocks: --> 124 x = module(x) 125 if isinstance(module, Conv_Block): 126 encoder_features.append(x) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/nn/modules/module.py:1736, in Module._wrapped_call_impl(self, *args, **kwargs) 1734 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc] 1735 else: -> 1736 return self._call_impl(*args, **kwargs) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/nn/modules/module.py:1747, in Module._call_impl(self, *args, **kwargs) 1742 # If we don't have any hooks, we want to skip the rest of the logic in 1743 # this function, and just call forward. 1744 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks 1745 or _global_backward_pre_hooks or _global_backward_hooks 1746 or _global_forward_hooks or _global_forward_pre_hooks): -> 1747 return forward_call(*args, **kwargs) 1749 result = None 1750 called_always_called_hooks = set() File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/nn/modules/pooling.py:296, in MaxPool3d.forward(self, input) 295 def forward(self, input: Tensor): --> 296 return F.max_pool3d( 297 input, 298 self.kernel_size, 299 self.stride, 300 self.padding, 301 self.dilation, 302 ceil_mode=self.ceil_mode, 303 return_indices=self.return_indices, 304 ) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/_jit_internal.py:624, in boolean_dispatch.<locals>.fn(*args, **kwargs) 622 return if_true(*args, **kwargs) 623 else: --> 624 return if_false(*args, **kwargs) File ~/miniconda3/envs/careamics/lib/python3.10/site-packages/torch/nn/functional.py:920, in _max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode, return_indices) 918 if stride is None: 919 stride = torch.jit.annotate(List[int], []) --> 920 return torch.max_pool3d(input, kernel_size, stride, padding, dilation, ceil_mode) NotImplementedError: The operator 'aten::max_pool3d_with_indices' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS." }System
Important
Environment
The text was updated successfully, but these errors were encountered: