Add option to use residual blocks for updownsampling #176

Warvito · 2023-01-10T11:53:05Z

Implements #173

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

…downsampling # Conflicts: # generative/networks/nets/diffusion_model_unet.py

marksgraham

I'm getting a torchscript error when I run the tests

marksgraham · 2023-01-11T21:33:49Z

tests/test_diffusion_model_unet.py

+    def test_script_conditioned_2d_models_with_resblock_updown(self):
+        net = DiffusionModelUNet(
+            spatial_dims=2,
+            in_channels=1,
+            out_channels=1,
+            num_res_blocks=1,
+            num_channels=(8, 8, 8),
+            attention_levels=(False, False, True),
+            norm_num_groups=8,
+            with_conditioning=True,
+            transformer_num_layers=1,
+            cross_attention_dim=3,
+            resblock_updown=True,
+        )
+        test_script_save(net, torch.rand((1, 1, 16, 16)), torch.randint(0, 1000, (1,)).long(), torch.rand((1, 1, 3)))


I'm getting errors on all the torchscript tests. I've tried with torch 1.13, 1.9 and 1.8:

forward(__torch__.generative.networks.nets.diffusion_model_unet.Downsample self, Tensor x) -> (Tensor): Expected at most 2 arguments but found 3 positional arguments. : File "/home/mark/projects/monai/GenerativeModels/generative/networks/nets/diffusion_model_unet.py", line 720 if self.downsampler is not None: if self.resblock_updown: hidden_states = self.downsampler(hidden_states, temb) ~~~~~~~~~~~~~~~~ <--- HERE else: hidden_states = self.downsampler(hidden_states)

It seems like it doesn't like that self.downsample can take a variable number of arguments depending on the code path. One hacky solution would be to allow the Downsampler and Upsampler forward methods to accept temb too, even if they don't use it

That is weird, it makes sense the error that you got but I cannot replicate it here. I am trying the tests with:

>>> from monai.config import print_config >>> print_config() MONAI version: 1.1.0+20.g7eddd2d8 Numpy version: 1.24.1 Pytorch version: 1.8.0+cu111 MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False MONAI rev id: 7eddd2d8f601471a0587ee9ae71c7325eca24ae7 MONAI __file__: /media/walter/Storage/Projects/GenerativeModels/MONAI/monai/__init__.py Optional dependencies: Pytorch Ignite version: 0.4.10 Nibabel version: 4.0.2 scikit-image version: NOT INSTALLED or UNKNOWN VERSION. Pillow version: 9.4.0 Tensorboard version: 2.11.0 gdown version: 4.6.0 TorchVision version: 0.9.0+cu111 tqdm version: 4.64.1 lmdb version: NOT INSTALLED or UNKNOWN VERSION. psutil version: 5.9.4 pandas version: NOT INSTALLED or UNKNOWN VERSION. einops version: 0.6.0 transformers version: NOT INSTALLED or UNKNOWN VERSION. mlflow version: NOT INSTALLED or UNKNOWN VERSION. pynrrd version: NOT INSTALLED or UNKNOWN VERSION. For details about installing the optional dependencies, please visit: https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies

and

>>> from monai.config import print_config >>> print_config() MONAI version: 1.1.0+20.g7eddd2d8 Numpy version: 1.24.0 Pytorch version: 1.13.1+cu117 MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False MONAI rev id: 7eddd2d8f601471a0587ee9ae71c7325eca24ae7 MONAI __file__: /media/walter/Storage/Projects/GenerativeModels/MONAI/monai/__init__.py Optional dependencies: Pytorch Ignite version: 0.4.10 Nibabel version: 4.0.2 scikit-image version: NOT INSTALLED or UNKNOWN VERSION. Pillow version: 9.3.0 Tensorboard version: 2.11.0 gdown version: 4.6.0 TorchVision version: 0.14.1+cu117 tqdm version: 4.64.1 lmdb version: NOT INSTALLED or UNKNOWN VERSION. psutil version: NOT INSTALLED or UNKNOWN VERSION. pandas version: NOT INSTALLED or UNKNOWN VERSION. einops version: 0.6.0 transformers version: NOT INSTALLED or UNKNOWN VERSION. mlflow version: NOT INSTALLED or UNKNOWN VERSION. pynrrd version: NOT INSTALLED or UNKNOWN VERSION. For details about installing the optional dependencies, please visit: https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies

I reinstalled my OS recently and might miss configured something. @marksgraham could you please share what your monai.config.print_config returns just to confirm it is not something related to the dependencies?
Here I started to get an error from the inferer tests after trying rerun all tests. Will add it in the issues

This is my monai config:

MONAI version: 1.1.dev2248 Numpy version: 1.23.4 Pytorch version: 1.8.0 MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False MONAI rev id: 3400bd91422ccba9ccc3aa2ffe7fecd4eb5596bf MONAI __file__: /home/mark/Envs/gen2/lib/python3.8/site-packages/monai/__init__.py Optional dependencies: Pytorch Ignite version: 0.4.10 Nibabel version: 5.0.0 scikit-image version: NOT INSTALLED or UNKNOWN VERSION. Pillow version: 9.3.0 Tensorboard version: 2.11.0 gdown version: 4.6.0 TorchVision version: 0.8.0 tqdm version: 4.64.1 lmdb version: NOT INSTALLED or UNKNOWN VERSION. psutil version: 5.9.4 pandas version: NOT INSTALLED or UNKNOWN VERSION. einops version: 0.6.0 transformers version: NOT INSTALLED or UNKNOWN VERSION. mlflow version: NOT INSTALLED or UNKNOWN VERSION. pynrrd version: NOT INSTALLED or UNKNOWN VERSION. For details about installing the optional dependencies, please visit: https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies

…downsampling

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Warvito · 2023-01-13T22:49:14Z

Sorry @marksgraham , I could not replicate the error here. I tried to make the changes following your recommendations and error description. Could you please check if they are okay?

Now I wonder if these errors that we are getting with the torchscript would be something that would appear in the MONAI's CI containers (since they are not appear in my new environment, and in theory I am using the same requirements). I will open an issue and try to check with @ericspod if there is any way that we could build or pull the testing docker images so we can run our tests locally.

marksgraham

Hi Walter, most the tests now pass but I'm getting one final, baffling test error:

ERROR: test_script_conditioned_2d_models_with_resblock_updown (__main__.TestDiffusionModelUNet2D)

.... 


RuntimeError:
Module 'Upsample' has no attribute 'conv' :
  File "/home/mark/projects/monai/GenerativeModels/generative/networks/nets/diffusion_model_unet.py", line 522
        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
                ~~~~~~~~~ <--- HERE
        return x

It can be fixed by removing the if statement on L506:

        if use_conv:
            self.conv = Convolution(
                spatial_dims=spatial_dims,
                in_channels=self.num_channels,
                out_channels=self.out_channels,
                strides=1,
                kernel_size=3,
                padding=padding,
                conv_only=True,

so that self.conv is always present, but I can't see how this error is being thrown unless the self.conv attribute is being set to True externally after the class is initalised, which I don't see happening anywhere.

marksgraham · 2023-01-13T23:01:15Z

Sorry @marksgraham , I could not replicate the error here. I tried to make the changes following your recommendations and error description. Could you please check if they are okay?

Now I wonder if these errors that we are getting with the torchscript would be something that would appear in the MONAI's CI containers (since they are not appear in my new environment, and in theory I am using the same requirements). I will open an issue and try to check with @ericspod if there is any way that we could build or pull the testing docker images so we can run our tests locally.

This is a good idea!

Warvito · 2023-01-13T23:03:02Z

I created #185 to discuss it

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Warvito · 2023-01-13T23:10:53Z

Yeah it latest error look like what we had in

GenerativeModels/generative/networks/nets/diffusion_model_unet.py

Line 691 in e49b35a

self.downsampler = None

I hope having else: self.conv = None might solve it

marksgraham · 2023-01-13T23:32:46Z

I still get an error - even though self.use_conv evaluates to False.... I've no idea what is going on here

cannot call a value of type 'NoneType':
  File "/home/mark/projects/monai/GenerativeModels/generative/networks/nets/diffusion_model_unet.py", line 526
        print('here', self.use_conv, self.use_conv is True)
        if self.use_conv:
            x = self.conv(x)
                ~~~~~~~~~ <--- HERE
        return x

Is this blocking you for other stuff? I'm starting to wonder if it is something wrong in my system causing this error - it doesn't make any sense. We could pass for now and raise an issue for further investigation - I can look into properly next week?

Warvito · 2023-01-13T23:43:44Z

Sounds good. I will create an issue for this error for future further investigation.
I have been mainly trying to get the issues related to an updated version of the diffusion model done before moving to the transformers, since it is involved in the anomaly detection and MRI reconstruction tutorials that the others are working on.

ericspod · 2023-01-14T22:16:54Z

@marksgraham @Warvito Torchscript is a little touchy on arguments, it won't like different number of arguments or overriding methods and changing signatures in certain ways. This is to maintain semantic compatibility with C++ as far as I can deduce, which is more strict in what can be done with methods and arguments.

Warvito added 2 commits January 10, 2023 11:48

Add residual blocks for down/upsampling

0f1df93

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Merge branch 'main' into 173-add-option-to-use-residual-blocks-for-up…

904f266

…downsampling # Conflicts: # generative/networks/nets/diffusion_model_unet.py

Warvito linked an issue Jan 10, 2023 that may be closed by this pull request

Add option to use residual blocks for up/downsampling. #173

Closed

Warvito marked this pull request as ready for review January 10, 2023 22:58

marksgraham self-requested a review January 11, 2023 21:20

marksgraham requested changes Jan 11, 2023

View reviewed changes

Warvito added 3 commits January 13, 2023 22:15

Merge branch 'main' into 173-add-option-to-use-residual-blocks-for-up…

d955afc

…downsampling

Add emb to Upsample and Downsample layers

374eccd

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Change calls #176

bdef518

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Warvito requested a review from marksgraham January 13, 2023 22:43

marksgraham requested changes Jan 13, 2023

View reviewed changes

Try to fix torchscript error

16989f6

Signed-off-by: Walter Hugo Lopez Pinaya <ianonimato@hotmail.com>

Warvito requested a review from marksgraham January 13, 2023 23:11

marksgraham approved these changes Jan 14, 2023

View reviewed changes

marksgraham merged commit 7dcdc9a into main Jan 14, 2023

Warvito deleted the 173-add-option-to-use-residual-blocks-for-updownsampling branch January 14, 2023 14:39

marksgraham mentioned this pull request Jan 16, 2023

Fix torchscript errors #192

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add option to use residual blocks for updownsampling #176

Add option to use residual blocks for updownsampling #176

Warvito commented Jan 10, 2023

marksgraham left a comment

marksgraham Jan 11, 2023

Warvito Jan 11, 2023 •

edited

marksgraham Jan 12, 2023

Warvito commented Jan 13, 2023

marksgraham left a comment

marksgraham commented Jan 13, 2023

Warvito commented Jan 13, 2023

Warvito commented Jan 13, 2023

marksgraham commented Jan 13, 2023

Warvito commented Jan 13, 2023

ericspod commented Jan 14, 2023

Add option to use residual blocks for updownsampling #176

Add option to use residual blocks for updownsampling #176

Conversation

Warvito commented Jan 10, 2023

marksgraham left a comment

Choose a reason for hiding this comment

marksgraham Jan 11, 2023

Choose a reason for hiding this comment

Warvito Jan 11, 2023 • edited

Choose a reason for hiding this comment

marksgraham Jan 12, 2023

Choose a reason for hiding this comment

Warvito commented Jan 13, 2023

marksgraham left a comment

Choose a reason for hiding this comment

marksgraham commented Jan 13, 2023

Warvito commented Jan 13, 2023

Warvito commented Jan 13, 2023

marksgraham commented Jan 13, 2023

Warvito commented Jan 13, 2023

ericspod commented Jan 14, 2023

Warvito Jan 11, 2023 •

edited