search

usnistgov / alignn

Atomistic Line Graph Neural Network https://scholar.google.com/citations?user=9Q-tNnwAAAAJ&hl=en
https://jarvis.nist.gov/jalignn/
Other
214 stars 79 forks source link

How to training the model with multiple CPUs? #164

Open LTJer opened 1 month ago

LTJer commented 1 month ago

Hi

How can I assign multiple cpus to the training process?

Thanks a lot.

bdecost commented 2 weeks ago

can you please clarify if you are asking about multithreaded training with the model on CPU vs using multiple dataloader processes when training with the model on GPU?

from what I've read online PyTorch defaults to one thread per physical core for parallel CPU ops. you can check the output of torch.get_num_threads to see if it's what you expect - if you're in an HPC environment you may want to double check what resources you are requesting

for multiple dataloader processes, the alignn config key num_workers sets the num_workers argument of the pytorch DataLoader