8029 update load old weights function for diffusion_model_unet.py

[yiheng-wang-nv]

Fixes #8029 .

Description

A few sentences describing the changes proposed in this pull request.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[yiheng-wang-nv]

Tested with this branch.
Code:

from generative.networks.nets import DiffusionModelUNet
from monai.networks.nets.diffusion_model_unet import DiffusionModelUNet as MonaiDiffusionModelUNet
import torch

input_params = {
    "spatial_dims": 2,
    "in_channels": 1,
    "out_channels": 1,
    "num_channels": [32, 64, 128, 256],
    "attention_levels": [False, True, True, True],
    "num_head_channels": [0, 32, 32, 32],
    "num_res_blocks": 2,
}

old_network = DiffusionModelUNet(**input_params)
old_network.load_state_dict(torch.load("model.pt"))

new_params = input_params.copy()
new_params.pop("num_channels")
new_params["channels"] = input_params["num_channels"]
new_params["include_fc"] = False
new_params["use_combined_linear"] = False

new_network = MonaiDiffusionModelUNet(**new_params)
new_network.load_old_state_dict(old_network.state_dict())

test_input = torch.randn(1, 1, 64, 64).cuda()
timesteps = torch.randn(1).cuda()

old_network = old_network.cuda()
new_network = new_network.cuda()

old_network.eval()
new_network.eval()

with torch.no_grad():
    old_output = old_network(test_input, timesteps)
    new_output = new_network(test_input, timesteps)

Output:

[KumoLiu]

Hi @yiheng-wang-nv, thanks for the update, could you please test https://github.com/Project-MONAI/MONAI/blob/dev/tests/test_diffusion_model_unet.py#L581? It looks like different branches lead to different blocks, and this test case doesn't seem to be able to be imported correctly in this pr.

[yiheng-wang-nv]

Thanks @KumoLiu . I updated the code to support both cases

[KumoLiu]

Thanks for the quick update, LGTM.

[KumoLiu]

Hi @ericspod, could you please also help review this one? Thanks!

[KumoLiu]

/build

[ericspod]

This looks fine to me as it's passing your tests.