Isonet Denoising Stuck at Iteration1

[AriaGarrett]

I'm trying to denoise tomograms using isonet in v0.7.2 using mostly default settings besides adjusting the SNR falloff and patch size. Preprocessing is always fine, and the denoising seems to work until the refine/first iteration begins. I've tried both with and without an initial model added but get the same error. Any thoughts?

2025-08-26 17:03:14 [INFO] pyp/system/local_run.py:77 | ######Isonet starts refining######
77 2025-08-26 17:03:18 [INFO] pyp/system/local_run.py:77 | 2025-08-26 17:03:18.112638: E tensorflow/stream_executor/cuda/cuda_blas.cc:2981] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
78 2025-08-26 17:03:23 [INFO] pyp/system/local_run.py:77 | 08-26 17:03:23, INFO Start Iteration1!
79 2025-08-26 17:03:24 [INFO] pyp/system/local_run.py:77 | 08-26 17:03:23, INFO Use Pretrained model as the output model of iteration 0 and predict subtomograms
80 2025-08-26 17:03:24 [INFO] pyp/system/local_run.py:77 | 08-26 17:03:23, INFO Start predicting subtomograms!
81 2025-08-26 17:09:35 [INFO] pyp/system/local_run.py:77 | 08-26 17:09:35, INFO Done predicting subtomograms!
82 2025-08-26 17:09:35 [INFO] pyp/system/local_run.py:77 | 08-26 17:09:35, INFO Noise Level:0.0
83 2025-08-26 17:13:51 [INFO] pyp/system/local_run.py:77 | 08-26 17:13:51, INFO Done preparing subtomograms!
84 2025-08-26 17:13:51 [INFO] pyp/system/local_run.py:77 | 08-26 17:13:51, INFO Start training!
85 2025-08-26 17:13:53 [INFO] pyp/system/local_run.py:77 | 08-26 17:13:53, INFO Loaded model from disk
86 2025-08-26 17:13:53 [INFO] pyp/system/local_run.py:77 | 08-26 17:13:53, INFO begin fitting
87 2025-08-26 17:13:53 [INFO] pyp/system/local_run.py:77 | Epoch 1/10
88 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | 08-26 17:14:09, ERROR Traceback (most recent call last):
89 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/pyp/external/IsoNet/bin/refine.py", line 128, in run
90 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | history = train_data(args) #train based on init model and save new one as model_iter{num_iter}.h5
91 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/pyp/external/IsoNet/models/unet/train.py", line 93, in train_data
92 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | history = train3D_continue('{}/model_iter{:0>2d}.h5'.format(settings.result_dir,settings.iter_count),
93 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/pyp/external/IsoNet/models/unet/train.py", line 64, in train3D_continue
94 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | history = model.fit(train_data, validation_data=test_data,
95 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
96 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | raise e.with_traceback(filtered_tb) from None
97 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/tensorflow/python/eager/execute.py", line 54, in quick_execute
98 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
99 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | tensorflow.python.framework.errors_impl.ResourceExhaustedError: Graph execution error:
100 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | Detected at node 'gradient_tape/model/fullconv_out/Conv3D/Conv3DBackpropInputV2' defined at (most recent call last):
101 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/pyp/external/IsoNet/bin/isonet.py", line 497, in
102 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | fire.Fire(ISONET)
103 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/fire/core.py", line 135, in Fire
104 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | component_trace = _Fire(component, args, parsed_flag_args, context, name)
105 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/fire/core.py", line 468, in _Fire
106 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | component, remaining_args = _CallAndUpdateTrace(
107 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/fire/core.py", line 684, in _CallAndUpdateTrace
108 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | component = fn(*varargs, **kwargs)
109 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/pyp/external/IsoNet/bin/isonet.py", line 391, in refine
110 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | run(d_args)
111 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/pyp/external/IsoNet/bin/refine.py", line 128, in run
112 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | history = train_data(args) #train based on init model and save new one as model_iter{num_iter}.h5
113 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/pyp/external/IsoNet/models/unet/train.py", line 93, in train_data
114 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | history = train3D_continue('{}/model_iter{:0>2d}.h5'.format(settings.result_dir,settings.iter_count),
115 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/pyp/external/IsoNet/models/unet/train.py", line 64, in train3D_continue
116 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | history = model.fit(train_data, validation_data=test_data,
117 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 65, in error_handler
118 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | return fn(*args, **kwargs)
119 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/engine/training.py", line 1564, in fit
120 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | tmp_logs = self.train_function(iterator)
121 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/engine/training.py", line 1160, in train_function
122 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | return step_function(self, iterator)
123 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/engine/training.py", line 1146, in step_function
124 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | outputs = model.distribute_strategy.run(run_step, args=(data,))
125 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/engine/training.py", line 1135, in run_step
126 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | outputs = model.train_step(data)
127 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/engine/training.py", line 997, in train_step
128 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
129 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/optimizers/optimizer_v2/optimizer_v2.py", line 576, in minimize
130 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | grads_and_vars = self._compute_gradients(
131 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/optimizers/optimizer_v2/optimizer_v2.py", line 634, in _compute_gradients
132 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | grads_and_vars = self._get_gradients(
133 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | File "/opt/conda/envs/isonet/lib/python3.9/site-packages/keras/optimizers/optimizer_v2/optimizer_v2.py", line 510, in _get_gradients
134 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | grads = tape.gradient(loss, var_list, grad_loss)
135 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | Node: 'gradient_tape/model/fullconv_out/Conv3D/Conv3DBackpropInputV2'
136 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | OOM when allocating tensor with shape[4,64,64,64,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
137 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | [[{{node gradient_tape/model/fullconv_out/Conv3D/Conv3DBackpropInputV2}}]]
138 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info. This isn't available when running in Eager mode.
139 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | [Op:__inference_train_function_31167]
140 2025-08-26 17:14:31 [INFO] /opt/pyp/bin/run/pyp:5144 | nextPYP (isonet train) finished successfully

[abartesaghi]

It seems like isonet is running out of memory in your system:

136 2025-08-26 17:14:10 [INFO] pyp/system/local_run.py:77 | OOM when allocating tensor with shape[4,64,64,64,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc

This type of error is usually solved by reducing the batch-size during training. However, in version 0.7.2 of nextpyp, a bug prevents this parameter from being exposed in the UI. Instead, you can try reducing the Cube size for training to 56 (or a smaller multiple of 8) and see if that helps (this may affect training performance though).

The issue with the batch-size parameter is trivial to fix, so I just went ahead and patched 0.7.2. If you reinstall that version, you should be able to reduce the Mini-batch size from the default value of 4 to something smaller.

[AriaGarrett]

Thank you! Yes, that fixed the issue.