implement group-norm - #790
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ba3f6ddeec
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| BlockType = Literal["conv_next_block", "conv_block"] | ||
| ActivationType = Literal["relu", "gelu", "capped_gelu"] | ||
| NormType = Literal["batch", "instance", "layer"] | ||
| NormType = Literal["batch", "instance", "group", "nonorm", "layer"] |
There was a problem hiding this comment.
Do not accept group/no-norm for ConvBlock until wired
When a config sets core_block.block_type: "conv_block" with the newly accepted norm: "group" or "nonorm", validation now succeeds, but the BlockConfig.build conv_block branch still constructs ConvBlock without passing self.norm, and ConvBlock hard-codes BatchNorm2d. That silently runs experiments with batch norm despite the YAML requesting group/no normalization, so either reject these values for conv_block or thread the norm setting through that block too.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
See my 🐑 suggestion above about the NormBuilder -- that idea could help address this concern.
alxmrs
left a comment
There was a problem hiding this comment.
A few ideas that could improve this group norm PR.
| self.checkpoint_simple = checkpoint_simple | ||
|
|
||
| @staticmethod | ||
| def _build_norm_layer( |
There was a problem hiding this comment.
It looks like this doesn't use any class members, therefore I think it could be a pure function not in the class rather than a method. Further, I think this could be generally useful -- how about we make it a public function?
|
|
||
| @staticmethod | ||
| def _build_norm_layer( | ||
| norm: str, channels: int, group_norm_groups: int |
There was a problem hiding this comment.
I think the parameter could just be groups?
🐑 maybe it could be an optional param since it's only used with group norm.
| self.checkpoint_simple = checkpoint_simple | ||
|
|
||
| @staticmethod | ||
| def _build_norm_layer( |
There was a problem hiding this comment.
🐑 idea: potentially, we could pass in to the core block a "NormBuilder" so we could set all of the hparams for which norm to use in our config system -- see the CoreBlockBuilder(Protocol) at the bottom of this file.
| BlockType = Literal["conv_next_block", "conv_block"] | ||
| ActivationType = Literal["relu", "gelu", "capped_gelu"] | ||
| NormType = Literal["batch", "instance", "layer"] | ||
| NormType = Literal["batch", "instance", "group", "nonorm", "layer"] |
There was a problem hiding this comment.
See my 🐑 suggestion above about the NormBuilder -- that idea could help address this concern.
| width=st.integers(min_value=1, max_value=16), | ||
| ) | ||
| @settings(max_examples=20) | ||
| @settings(max_examples=20, deadline=None) |
There was a problem hiding this comment.
I think we should remove this -- if there's a timeout in our property based tests, that's a concern worth knowing about.
implement group and layer normalization with tests. Change norm selection to use build_norm_layer that is more extensible than old if/elif logic.