search

hpcaitech / FastFold

Optimizing AlphaFold Training and Inference on GPU Clusters
Apache License 2.0
557 stars 86 forks source link

support training, fix some inplace func of nn #118

Closed Gy-Lu closed 1 year ago

Gy-Lu commented 1 year ago

This PR should be merged after the other two.

  1. Add training code and its shell script, support DDP now.
  2. Add lr_scheduler and move loss to hub.
  3. Modify some bugs introduced by openfold, the inplace func in dropout and the unused out_product_mean in evoformer.

For unit test, it would cause OOM on CI machine. I am trying to develop a tiny model.

Shenggan commented 1 year ago

Seem still have some bugs:

  1. from fastfold.model.loss import * -> from fastfold.model.hub.loss import * in fastfold/model/nn/heads.py
  2. https://github.com/hpcaitech/FastFold/blob/164f67773f971b57a213d0f402e58e6666fb8a68/fastfold/model/fastnn/template.py#L394 introduced by https://github.com/hpcaitech/FastFold/pull/117