Refactor Attention UNet #115
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
1. Removed DoubleConvBlock which was being used for the residual connection. Instead a skip connection is directly included in the FeatureIntegrationModule. 2. Time step embedding (and other embeddings) are now added within the residual block, rather than at the end. 3. Removed some warnings. 4. Added `num_attention_heads` parameter to control the self attention and cross attention heads in the model
HarshithBachimanchi
requested review from
JesusPinedaC,
BenjaminMidtvedt and
giovannivolpe
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR aims to refactor
AttentionUNet. The following changes are made to make it modularly simple, and produce a much cleaner print statement:DoubleConvBlockis now removed. It was used to implement residual connections at all levels in the UNet. Instead, now they are directly integrated to theFeatureIntegrationModulethrough a couple ofBlockmodules.num_attention_headsto control the number of attention heads in self-attention and cross-attention heads. Also added a unittest for this.To refactor into deeply style, I see that several styles need to be implemented for
Conv2dBlock, followed by their integration toUNet2d. Some of them exist, but not in the way I want (For example styles,spatial_self_attention, andspatial_cross_attention). I gave it a quick try, and looks like it is possible to build it with styles but requires extensive testing. I prefer to keepAttentionUNetbespoke for now (Unless if you have any suggestions).