usnistgov / alignn

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

Pytorch version of the code #59

Open kdmsit opened 2 years ago

kdmsit commented 2 years ago

Is it possible to have the torch version of the code in place of torch-ignite? I want to go deep dive into the actual code and want to make some changes in loss function. I have tried to replace the trainer.run(train_loader, max_epochs=config.epochs) with simple iterative torch version of it as follows :

for epoch in range(config.epochs): for batch in train_loader: net = net.train() optimizer.zero_grad() graph, line_graph, target = batch if torch.cuda.is_available(): graph = graph.to(device, non_blocking=True) line_graph = line_graph.to(device, non_blocking=True) target = target.to(device, non_blocking=True) g=(graph,line_graph) output = net(g) loss = criterion(output, target) mae_error = mae(output.data.cpu(), target.cpu()) loss.backward() optimizer.step()

But I am not able to reproduce the result. Could you please help me to resolve the issue?

bdecost commented 2 years ago

Hi @kdmsit -- would you mind giving some more details about the loss function you're using and what is not working as expected?

One thing that I noticed is that you aren't using a learning rate scheduler in this example -- if you want to explicitly write the training loop, the default scheduler we use is cosine annealing with linear warmup, and you'll want to call scheduler.step() on every iteration, just after optimizer.step(), in the inner loop.

If you just want a different (or custom) loss, you can add it here (and to the configuration settings) and keep using ignite. I've found that ignite really helps reduce the duplication and complexity of evaluation logic, for example.

Do you have a working example you could share maybe?