The tutorial accompanying the pull request #7308 for MONAI core, which adds SURE-loss and Conjugate Gradient #1631
Conversation
Signed-off-by: chaoliu <[email protected]>
Signed-off-by: chaoliu <[email protected]>
Signed-off-by: chaoliu <[email protected]>
Signed-off-by: chaoliu <[email protected]>
Signed-off-by: chaoliu <[email protected]>
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
for more information, see https://pre-commit.ci
generative/smrd/models/models.py
Outdated
| from .normalization import get_normalization | ||
|
|
||
|
|
||
| class UNet(nn.Module): |
There was a problem hiding this comment.
I think we may need to use the components from the MONAI repo.
There was a problem hiding this comment.
Did you mean this one? We defined the UNet ourselves here since the pretrained model we are using is from this definition. Otherwise, we need to convert the weights/names to be consistent to the Unet from MONAI.
There was a problem hiding this comment.
Yes, do we have specific reason to add this pretrained model otherwise we can just ensure functionality level works?
What do you think? cc @ericspod
There was a problem hiding this comment.
It would be best wherever possible to use MONAI components. We have our own UNet class but is structured differently from this one, and we do have other blocks and layers that could be substituted for those defined here. If possible please use those in MONAI unless it causes significant headaches converting the weight names in your pretrained data.
generative/smrd/SMRD.ipynb
Outdated
There was a problem hiding this comment.
I would put the code for SMRDOptimizer into its own file since it's so large in this notebook.
generative/smrd/models/models.py
Outdated
|
|
||
| self.end_conv = nn.Conv2d(ngf, config.data.channels, 3, stride=1, padding=1) | ||
|
|
||
| self.res1 = nn.ModuleList( |
There was a problem hiding this comment.
You could use Sequential instead of ModuleList here and elsewhere in this class, then you wouldn't need _compute_cond_module.
| class NoneNorm2d(nn.Module): | ||
| def __init__(self, num_features, bias=True): | ||
| super().__init__() | ||
|
|
||
| def forward(self, x): | ||
| return x |
There was a problem hiding this comment.
You can actually just use torch.nn.Identity whose constructor takes any arguments and just does what this class does.
|
View / edit / reply to this conversation on ReviewNB ericspod commented on 2024-02-07T13:03:59Z Line #1. def denoise_cg_step( Would it make sense to put this in Core with the loss function definition? |
|
View / edit / reply to this conversation on ReviewNB ericspod commented on 2024-02-07T13:04:00Z The key thing in this tutorial is what the optimizer is doing with the loss function to do the reconstruction. I know it's explained in the paper but I would much rather have more explanation here with more insightful comments in the code of the optimizer. I have mentioned in other comments about simplifying the class by removing unneeded definitions and moving utility methods elsewhere, making the class simpler will help in understanding what it's doing. You can also link/cite to the paper here to point people in the right direction. |
|
View / edit / reply to this conversation on ReviewNB ericspod commented on 2024-02-07T13:04:00Z Line #23. def _dict2namespace(self, config): Methods like this which don't use |
|
View / edit / reply to this conversation on ReviewNB ericspod commented on 2024-02-07T13:04:01Z Line #60. return MulticoilForwardMRI(self.config["orientation"])( I would suggest creating an object instance of MulticoilForwardMRI in the constructor and use that in this function. This avoids creating the object every time the function is called. |
Looks good to me, I think this stuff is complementary to the diffusion work we're doing. As already mentioned by others, would have been great to use the monai networks already defined for the pretrained network (would cut down on the code added and make it easier for others to integrate this into their own work). I wonder if it would be hard to re-train the network using a MONAI implementation and then integrate that? Might be easier than trying to convert weights over |
Co-authored-by: Eric Kerfoot <[email protected]> Signed-off-by: cxlcl <[email protected]>
Co-authored-by: Eric Kerfoot <[email protected]> Signed-off-by: cxlcl <[email protected]>
Co-authored-by: Eric Kerfoot <[email protected]> Signed-off-by: cxlcl <[email protected]>
Co-authored-by: Eric Kerfoot <[email protected]> Signed-off-by: cxlcl <[email protected]>
Signed-off-by: chaoliu <[email protected]>
for more information, see https://pre-commit.ci
|
Thanks a lot for all the suggestions. Based on the suggestions, I've made several changes:
Please let me know how it looks and thanks again for providing the feedbacks! |
| class EMAHelper: | ||
| def __init__(self, mu=0.999): | ||
| self.mu = mu | ||
| self.shadow = {} | ||
|
|
||
| def register(self, module): | ||
| if isinstance(module, nn.DataParallel): | ||
| module = module.module | ||
| for name, param in module.named_parameters(): | ||
| if param.requires_grad: | ||
| self.shadow[name] = param.data.clone() | ||
|
|
||
| def update(self, module): | ||
| if isinstance(module, nn.DataParallel): | ||
| module = module.module | ||
| for name, param in module.named_parameters(): | ||
| if param.requires_grad: | ||
| self.shadow[name].data = (1.0 - self.mu) * param.data + self.mu * self.shadow[name].data | ||
|
|
||
| def ema(self, module): | ||
| if isinstance(module, nn.DataParallel): | ||
| module = module.module | ||
| for name, param in module.named_parameters(): | ||
| if param.requires_grad: | ||
| param.data.copy_(self.shadow[name].data) |
There was a problem hiding this comment.
A suggested shortening:
| class EMAHelper: | |
| def __init__(self, mu=0.999): | |
| self.mu = mu | |
| self.shadow = {} | |
| def register(self, module): | |
| if isinstance(module, nn.DataParallel): | |
| module = module.module | |
| for name, param in module.named_parameters(): | |
| if param.requires_grad: | |
| self.shadow[name] = param.data.clone() | |
| def update(self, module): | |
| if isinstance(module, nn.DataParallel): | |
| module = module.module | |
| for name, param in module.named_parameters(): | |
| if param.requires_grad: | |
| self.shadow[name].data = (1.0 - self.mu) * param.data + self.mu * self.shadow[name].data | |
| def ema(self, module): | |
| if isinstance(module, nn.DataParallel): | |
| module = module.module | |
| for name, param in module.named_parameters(): | |
| if param.requires_grad: | |
| param.data.copy_(self.shadow[name].data) | |
| class EMAHelper: | |
| def __init__(self, mu=0.999): | |
| self.mu = mu | |
| self.shadow = {} | |
| def _get_parameters(self,module): | |
| if isinstance(module, nn.DataParallel): | |
| module = module.module | |
| return module.named_parameters() | |
| def register(self, module): | |
| for name, param in self._get_parameters(module): | |
| if param.requires_grad: | |
| self.shadow[name] = param.data.clone() | |
| def update(self, module): | |
| for name, param in self._get_parameters(module): | |
| if param.requires_grad: | |
| self.shadow[name].data = (1.0 - self.mu) * param.data + self.mu * self.shadow[name].data | |
| def ema(self, module): | |
| for name, param in self._get_parameters(module): | |
| if param.requires_grad: | |
| param.data.copy_(self.shadow[name].data) |
| def ema_copy(self, module): | ||
| if isinstance(module, nn.DataParallel): | ||
| inner_module = module.module | ||
| module_copy = type(inner_module)(inner_module.config).to(inner_module.config.device) |
There was a problem hiding this comment.
Can we not do inner_module.clone()? This line will work if the type of the inner module is something constructable like this, but the Pytorch pattern is to use the clone method then the load_state_dict() method call isn't needed.
|
View / edit / reply to this conversation on ReviewNB ericspod commented on 2024-02-26T18:27:23Z Here I would suggest using download_and_extract since these commands will not work under Windows. You should be able to set things up correctly in Python code and not have to resort to command line calls. |
|
View / edit / reply to this conversation on ReviewNB ericspod commented on 2024-02-26T18:27:24Z This is great explanation. I would add a short paragraph just after the title at the top of this notebook echoing this here, describing what SMRD is and what the SURE loss is used for ("The SURE loss is used to decide when to stop the generation of the dense MRI, to avoid artifacts due to excessive sampling iterations." is a great summation of the motivation). |
|
This looks much better and much smaller. I had a few comments in the notebook and code but overall it's looking great. The CI fails can be fixed as well, it's too bad this can't be tested with papermill at the moment but we can live with it. |
|
View / edit / reply to this conversation on ReviewNB KumoLiu commented on 2024-02-27T08:12:24Z All of our tutorial notebooks are standardized to commence with these three sections: License, Setup Environment, and Setup Imports. To ensure that the ci checks pass successfully, could you please adjust your notebook to align with this structure? https://github.com/Project-MONAI/tutorials/actions/runs/8035965820/job/21949132998?pr=1631#step:7:10
You can refer here: https://github.com/Project-MONAI/tutorials/blob/main/CONTRIBUTING.md#create-a-notebook |
|
View / edit / reply to this conversation on ReviewNB KumoLiu commented on 2024-02-27T08:12:25Z Line #7. import pickle, gzip
I've noticed several discrepancies with the PEP8 style guide in the current work. Could you kindly address these to ensure our codebase maintains its consistency and readability? I appreciate your cooperation in adhering to our coding standards. |
|
View / edit / reply to this conversation on ReviewNB KumoLiu commented on 2024-02-27T08:12:25Z Line #15. from smrd_optimizer import SMRDOptimizer Where does this cxlcl commented on 2024-02-28T00:21:13Z Added in the latest commit. Forgot to include it in the previous commit... Thanks for the comments for the code style. I will address it later this week. |
|
View / edit / reply to this conversation on ReviewNB KumoLiu commented on 2024-02-27T08:12:26Z Line #2. from monai.networks.layers import ConjugateGradient I also recommend moving these import parts into "Setup imports". |
|
View / edit / reply to this conversation on ReviewNB KumoLiu commented on 2024-02-27T08:12:27Z Line #2. from monai.losses.sure_loss import SURELoss Also here. |
Signed-off-by: chaoliu <[email protected]>
for more information, see https://pre-commit.ci
|
Added in the latest commit. Forgot to include it in the previous commit... Thanks for the comments for the code style. I will address it later this week. View entire conversation on ReviewNB |
|
Hi @cxlcl sorry for losing track on this PR. Could you please fix the DCO issue and see if there's anything in the unresolved comments that need to be addressed still? Thanks! |
|
Hi @ericspod , sorry for having not attended to this for long. I will look into those DCO issue this week. As for the comments, I guess all the remaining comments and suggestions are about the DCO or pep8 standard. |
|
Hi @cxlcl if you're still wanting to finish this tutorial we can definitely review once the issues are sorted. Things have changed a bit so the conflicts and other things will have to be sorted before review and merging. Thanks! |
|
Hi Eric, Sorry for the late response. Yes, I would like to resolve the DCO issue. I've tried to resolve it by first run the local check script and upload. But the issue is still there in the CI . Is there any other way to check it locally or quickly without waiting for the results from the online CI process? |
The DCO details describes how to do an empty remedial commit. I don't see that here in the commit logs so I don't think you've pushed it if you tried. The other CI fails are for different things you have to resolve, as well as the conflicts. |
The tutorial accompanying the pull request #7308 for MONAI core, which adds SURE-loss and Conjugate Gradient
Description
A few sentences describing the changes proposed in this pull request.
Checks
./figurefolder./runner.sh -t <path to .ipynb file>