search

molecule-one / megan

Code for "Molecule Edit Graph Attention Network: Modeling Chemical Reactions as Sequences of Graph Edits"
MIT License
57 stars 19 forks source link

Speedup for training #8

Closed SongtaoLiu0823 closed 1 year ago

SongtaoLiu0823 commented 1 year ago

https://github.com/molecule-one/megan/blob/master/src/model/megan.py#L19

I think this line can be replaced by: if torch.cuda.is_available(): onehot = torch.cuda.FloatTensor(*x.shape, dims).zero() else: onehot = torch.FloatTensor(*x.shape, dims).zero().to(device)

As a result, the cpu usage can be reduced, leading to faster training.

SongtaoLiu0823 commented 1 year ago

In my machine, the CPU usage is from 2000+ to 100+.

mikolajsacha commented 1 year ago

Hi, sorry for the late response, and thank you for your find. We will merge it into the codebase :)