Issue using `decode_dicom_image` with `tf.vectorized_map`

[innat]

Backgrouond

I have a dataframe with dicom file path. I would like to create a keras layer, where it will read dicom file and do some preocessing (resizing , normalization etc). And return 2D input tensor. I am trying to add this layer with actual model to build a final model. The purpose is to get leverage the image processing on GPU.

Issue

In the call method of this layer, I tried to use tf.vectorized_map, but after running on some files, it gave an error. That is

InvalidArgumentError: PartialTensorShape: Incompatible shapes during merge: [1,5355,4915,1] vs. [1,2776,2082,1]:Decode DICOMI mage/for/TensorList ConcatV2}}]]

It might be due to the variable image array after decoding the dicom file. And tensorflow vectorized map might have an issue with concat function, as it has known limitation.

class DicomReadAndResizer(layers.Layer):
    def __init__(
        self, 
        input_height, 
        input_width, 
        **kwargs
    ):
        super().__init__(**kwargs)
        self.input_height = input_height
        self.input_width = input_width

    def read_and_resize_dicom(self, x):
        image_bytes = tf.io.read_file(x)
        image = tfio.image.decode_dicom_image(
            image_bytes, 
            dtype=tf.float32
        )
        image = tf.image.grayscale_to_rgb(image)
        image = tf.image.resize(
            image, 
            size=[self.input_height, self.input_width],
        )
        return image
     
    def call(self, inputs):
         # Works some-what
        x = tf.map_fn(
            self.read_and_resize_dicom, inputs, fn_output_signature=tf.float32,
        ) 

        # Doesn't work
        x = tf.vectorized_map(self.read_dicom, inputs)
        return x

Note, to make tf.map_fn to work, I have to place image resizer function and dicom reader in the same method (above self.read_and_resize), otherwise this tf.map_fn function also creates issue with multi-scale input tensor. But that doesn't work with tf.vectorized_map.

Questions

1. Could be there any fix in tf.io side? (maybe use tf.stack instead of tf.concat!) ? Is there any hacks that can be used to make vectorize function work?
2. Even if I am using such layer inside a model (as a hope to leverage GPU), will this layer still somehow swipe to CPU mode? Because, when I run the code (with tf.map_fn), the CPU processing still gets very high and memory consumption increases eventually.

About the data-loader, it is a very simple tf.data API. I pass a dataframe, where the path of dicom file is kept. So, there is not any heavy load for CPU side. Could you please give some pointer? What could be the reason? Thanks.

model = keras.Sequential(
            [
                keras.layers.InputLayer(dtype=tf.string),
                DicomReadAndResizer(
                    *INP_SIZE,
                ), # output [None, h, w, 3]
                trained_model # input [None, h, w, 3]
            ]
        )

[innat]

cc. @MarkDaoust

[MarkDaoust]

This is more like a stack overflow question than an issue with the repository.

I tried to use tf.vectorized_map
x = tf.vectorized_map(self.read_dicom, inputs)

vectorized_map is probably failing because your images have different sizes. You probably want read_and_resize there.

vectorized map might have an issue with concat function, as it has known limitation.

Concat's only failing because you're returning images of different sizes.

But that doesn't work with tf.vectorized_map.

Oh, what goes wrong? Can you give more details?

the CPU processing still gets very high and memory consumption increases eventually.

I doubt there's a GPU kernel for decode_dicom so that will happen on CPU regardless.

About the data-loader, it is a very simple tf.data API

It may be easier to work with this if instead of doing the transformations through tf.map_fun you use ds.map while you still have single images to load and resize. But note that tf.data usually runs everything on CPU to free up the GPU for training the model, which is usually the bottleneck.

[innat]

@MarkDaoust

vectorized_map is probably failing because your images have different sizes.

Yes, sort of. But I was also expecting, if it could handle such issue.

the CPU processing still gets very high and memory consumption increases eventually.
I doubt there's a GPU kernel for decode_dicom so that will happen on CPU regardless.

That is really unfortunate. I asked regarding this here. Hers is one inspiration https://developer.nvidia.com/dali

But note that tf.data usually runs everything on CPU to free up the GPU for training the model, which is usually the bottleneck.

Yes, that's why I wanted to just parse file from tf.data and wanted reading-resizing entirely on GPU. It is just for to get speed up.

But that doesn't work with tf.vectorized_map.
Oh, what goes wrong? Can you give more details?

I can definitely do that. I can share reproducible code. Here is the log.

InvalidArgumentError:  PartialTensorShape: Incompatible shapes during merge: [1,3328,2560,1] vs. [1,4096,3328,1]
	 [[{{node loop_body/IO>DecodeDICOMImage/pfor/TensorListConcatV2}}]] [Op:__inference_f_1285]

Function call stack:
f

[bhack]

That is really unfortunate. I asked regarding this here. Hers is one inspiration https://developer.nvidia.com/dali

We already talked about this at:
keras-team/keras-cv#146 (comment)

[MarkDaoust]

Thanks @bhack

PartialTensorShape: Incompatible shapes during merge: [1,3328,2560,1] vs. [1,4096,3328,1]

Yeah, you just need to do a resize first. It can't vectorize it if the outputs aren't the same size.

wanted reading-resizing entirely on GPU. It is just for to get speed up.

Are you sure that's the bottle-neck in your case. Usually the GPU is busy enough running the model that you have time to decode on CPU in the background.

[innat]

Are you sure that's the bottle-neck in your case. Usually the GPU is busy enough running the model that you have time to decode on CPU in the background.

Actually, I tired this for inference. (some interesting discussion, Nvidia DALI for decode-dicom, and dicomsdl > pydicom

PartialTensorShape: Incompatible shapes during merge: [1,3328,2560,1] vs. [1,4096,3328,1]
Yeah, you just need to do a resize first. It can't vectorize it if the outputs aren't the same size.

Could you please take a quick look to this colab? I did resize, which worked for tf.map_fn but not for tf.vectorized_map. If you want to test it with data, then please create a notebook with this kaggle data. Then import the colab code and just run all cell. No further change is required to reproduce the issue.