Training GCPN on a custom dataset

[akshat998]

Hello,
I was able to generate new molecules after training on a subset of ZINC molecules (as described in the tutorial):

import torch
from torchdrug import datasets

dataset = datasets.ZINC250k("~/molecule-datasets/", kekulize=True,
                            atom_feature="symbol")


from torchdrug import core, models, tasks

model = models.RGCN(input_dim=dataset.node_feature_dim,
                    num_relation=dataset.num_bond_type,
                    hidden_dims=[256, 256, 256, 256], batch_norm=False)
task = tasks.GCPNGeneration(model, dataset.atom_types, max_edge_unroll=12,
                            max_node=38, criterion="nll")

from torch import nn, optim
optimizer = optim.Adam(task.parameters(), lr = 1e-3)
solver = core.Engine(task, dataset, None, None, optimizer,
                     gpus=(0,), batch_size=128, log_interval=10)

solver.train(num_epoch=1)
solver.save("path_to_dump/graphgeneration/gcpn_zinc250k_1epoch.pkl")

I was wondering if it might be possible to pre-train the model on a custom dataset of molecules? :)

Thank you for the help!