Rank2 tensor head #792
Rank2 tensor head #792
Conversation
Codecov ReportAttention: Patch coverage is
|
|
I ran comparisons between the original rank2 stress tensor code that @lbluque has and compared it to the Hydra equivalent model. The results initially were different, however we traced the issue down to a difference in using amp in the head. This PR moves all heads to by default be used without AMP. |
| @@ -189,6 +189,10 @@ def no_weight_decay(self) -> list: | |||
|
|
|||
|
|
|||
| class HeadInterface(metaclass=ABCMeta): | |||
| @property | |||
| def use_amp(self): | |||
| return False | |||
There was a problem hiding this comment.
should the default be to not use it?
There was a problem hiding this comment.
When i talked to @wood-b he mentioned that we should have it off by default, it shouldnt change much memory usage and from our experiments it only helps with performance metrics
* rank 2 tensor head * fix rank2 head and add to e2e test * small fixes * keep hydra graphmixin * add rank2 head tests * test fixes * fix tests; move init_weight out of equiformer; add amp property to heads+hydra * add amp to heads and hydra * fix import * update snapshot; fix test seed and change tolerance --------- Co-authored-by: Misko <misko@meta.com> Former-commit-id: 6019dbce045d8ba26d711361197fad17d96ebe9b
* rank 2 tensor head * fix rank2 head and add to e2e test * small fixes * keep hydra graphmixin * add rank2 head tests * test fixes * fix tests; move init_weight out of equiformer; add amp property to heads+hydra * add amp to heads and hydra * fix import * update snapshot; fix test seed and change tolerance --------- Co-authored-by: Misko <misko@meta.com> Former-commit-id: 93afd03b796fd354cad5acd9a8676bee81a4a045
This PR adds a rank 2 tensor head for equiformerV2.
It is based on the implementations in the stress_trainer branch:
https://github.com/FAIR-Chem/fairchem/blob/stress_trainer/ocpmodels/models/utils/outer_block.py
TODO