Chunking issue in structure_tensor/multiprocessing.py

[chourroutm]

To prevent edge effects with large $\sigma$ or $\rho$, padding or overlap is often used, but I don't see this feature in structure_tensor/multiprocessing.py. Is it actually implemented in the "multiprocessing" package? If so, how is the overlap defined? By percents of block size or fixed amount of voxels or based on $\sigma$ or $\rho$?

I could find it in the legacy code hosted on Zenodo.

Dask has an implementation for this usecase (dask/array/overlap.py), but it is or used to be buggy.

[Skielex]

multiprocessing use util.get_block and util.insert_block functions to ensure correct overlap of the blocks based on rho and sigma. These are more or less updated versions of the legacy code you reference on Zenodo.

Have you experienced any issues with the code?

[chourroutm]

Yes, there is a chunking issue for sigma=3, rho=6, truncate=3 over a $1493\times16818\times64$ volume. (I don't know if the normalization issue across RGB is related.)

Here is the middle slice:

Overview of RGB slice

Overview and zoom of data slice

The code I used:

from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.image as im
import napari

import numpy as np
import scipy as sp
import cupy as cp
import numba
from dask.diagnostics import ProgressBar
import zarr
import dask
import dask.array
import tifffile as tiff

from structure_tensor.cp import structure_tensor_3d as structure_tensor_3d_gpu, eig_special_3d as eig_special_3d_gpu
from structure_tensor import structure_tensor_3d as structure_tensor_3d_cpu, eig_special_3d as eig_special_3d_cpu
from structure_tensor.multiprocessing import parallel_structure_tensor_analysis

data_path = Path("/hdd/output")
data_path.mkdir(exist_ok=True)

dataset_full = dask.array.from_zarr("gs://my_data.zarr") # dataset is private

lower_corner = (2325, 1113, 10774-32)
upper_corner = (17268, 17931, 10774+32)

optimal_intensity_window = (14246, 14456)

data_slab = dataset_full[
    lower_corner[0]:upper_corner[0], # NB: reverse order as the data is stored as z-y-x
    lower_corner[1]:upper_corner[1],
    lower_corner[2]:upper_corner[2]
].rechunk((512,512,64)).astype(np.float64)

filename = "data_preview.tif"
if not (data_path / Path(filename)).exists():
    print(f"Saving {filename}")
    tiff.imwrite(data_path / Path(filename),data_slab[:,:,32],photometric='minisblack')

data_slab[data_slab == 0] = np.nan
data_slab_inmem = data_slab.astype(np.float64).compute()

sigma = 3
rho   = 6
truncate = 3

print(f"Computing structure tensor and eigenvectors")
_, dirmin_slab_inmem, _ = parallel_structure_tensor_analysis(data_slab_inmem, sigma, rho, truncate=truncate, block_size=512, devices=32*['cpu']+ 16*['cuda:0'] + 16*['cuda:1'])

dirmin_slab_inmem[:, np.isnan(data_slab_inmem)] = np.nan

print(f"Saving vector field")
np.save(data_path / Path(f"vectorfield_st_sigma={sigma}_rho={rho}_truncate={truncate}.npy"),dirmin_slab_inmem)

print(f"Computing RGB")
rgb_slab_inmem = dirmin_slab_inmem / np.linalg.norm(dirmin_slab_inmem,axis=0) * 255
rgb_slab_inmem[:, np.any(np.isnan(rgb_slab_inmem),axis=0)] = 0

del dirmin_slab_inmem

rgb_slab_inmem = rgb_slab_inmem.astype(np.uint8)

print(f"Saving RGB preview as TIFF")
tiff.imwrite(data_path / Path(f"rgb_st_sigma={sigma}_rho={rho}_truncate={truncate}_preview.tif"),rgb_slab_inmem[[2,1,0],:,:,32].transpose((2,1,0)),photometric='rgb',imagej=True)

print(f"Saving RGB")
np.save(data_path / Path(f"rgb_st_sigma={sigma}_rho={rho}_truncate={truncate}.npy"),rgb_slab_inmem)

print("Done!")

[Skielex]

Based on the middle slice, it looks more like there's a difference between blocks. Try running exclusively on GPU or CPU insted of mixing.

parallel_structure_tensor_analysis(data_slab_inmem, sigma, rho, truncate=truncate, block_size=512, devices=16*['cuda:0'] + 16*['cuda:1'])

parallel_structure_tensor_analysis(data_slab_inmem, sigma, rho, truncate=truncate, block_size=512, devices=32*['cpu'])

That will help us pinpoint if the issue is related to the CPU or GPU implementation.

Side note: In most use cases I'd probably avoid mixing CPU and GPU computations, as it may introduce small run-to-run variations. CPU and GPU results may vary slightly as the order of operations in the gaussian_filter implementations may vary.

[chourroutm]

Thanks, giving it a test.

Also, as you may have noticed, I converted my lazy-loading dataset to full in-memory, but does parallel_structure_tensor_analysis also support lazy loading?

[chourroutm]

You were right! The processing was faster and without these different chunk issues. But shouldn't the computation output be uniformized between CPU and GPU?

Also, I still have an issue with normalizing where the blue channel has most of the dynamic (initially the red channel, but I did a swap with rgb_slab_inmem[[2,1,0],:,:,32].transpose((2,1,0))), any idea where it's from? Could it be due to the padding around the image? I would eventually change

structure-tensor/structure_tensor/util.py

Line 30 in 44c3ac1

kernel_radius = int(sigma * truncate + 0.5)

to:

    kernel_radius = int(max(sigma, rho) * truncate + 0.5)

[Skielex]

Thanks, giving it a test.

Also, as you may have noticed, I converted my lazy-loading dataset to full in-memory, but does parallel_structure_tensor_analysis also support lazy loading?

Only np.ndarray (and np.memmap) are currently supported by parallel_structure_tensor_analysis and I don't have plans to add support for anything else. The code is already quite complex considering how simple the task actually is, just because there are so many different options. Adding support for more data formats would make the function even more complicated and brittle.

I think it's probably better to implement your own version for your specific data type and hardware using the underlying get_block, structure_tensor_3d, eig_special_3d, and insert_block functions and a ProcessPoolExecutor or something. That way you can easily load data lazily. If dask provides good blocking functionality out of the box, you can use those instead of the structure_tensor.util functions. You could probably do it in less than 100 lines of Python.

You were right! The processing was faster and without these different chunk issues. But shouldn't the computation output be uniformized between CPU and GPU?

Yes, the output should be the same (save for numerical rounding errors), but clearly something is wrong (hence the brittleness comment). Which one is producing bad results, CPU or GPU?

Also, I still have an issue with normalizing where the blue channel has most of the dynamic (initially the red channel, but I did a swap with rgb_slab_inmem[[2,1,0],:,:,32].transpose((2,1,0))), any idea where it's from? Could it be due to the padding around the image? I would eventually change

I'm not sure what you mean by color channels. The input is greyscale volume and the output is eigenvalues and eigenvectors for each voxel.

structure-tensor/structure_tensor/util.py

Line 30 in 44c3ac1

kernel_radius = int(sigma * truncate + 0.5)
to:
kernel_radius = int(max(sigma, rho) * truncate + 0.5)

Right now the get_block function just takes one value. It's up to the caller to ensure this is max(sigma, rho).

structure-tensor/structure_tensor/multiprocessing.py

Lines 394 to 401 in 44c3ac1

	block, pos, pad = util.get_block(
	block_id,
	_data_sources.data,
	sigma=max(_worker_args.sigma, _worker_args.rho),
	block_size=_worker_args.block_size,
	truncate=_worker_args.truncate,
	copy=False,
	)

[Skielex]

@chourroutm, do you get similar results with GPU and CPU when you don't mix them?

[Skielex]

I'm running some tests and there appears to be an issue with the output when using pool_type="process".

[chourroutm]

@chourroutm, do you get similar results with GPU and CPU when you don't mix them?

I need to try!

I'm running some tests and there appears to be an issue with the output when using pool_type="process".

Okay, was it the default settings?

[Skielex]

@chourroutm , I haven't had time to further investigate the issue with differences between CPU and GPU results when you process-based pool.

I did fix an error that resulted in incorrect truncation of data 8d46c57, but that only resulted in very small differences when comparing the chunked results to plain results. I haven't published the fix to PyPI yet, as I was hoping to fix the other issue too. However, I'm not sure when I'll find time to actually work on this.

[chourroutm]

Okay, thanks for the update!