<jupyter_start><jupyter_text>Datamodule<jupyter_code>from torch_geometric.datasets import QM9

from shared import QM9MainDatamodule
from torch_geometric.loader import DataLoader as GeomDataLoader


class GraphQM9Datamodule(QM9MainDatamodule):
    def train_dataloader(self):
        return GeomDataLoader(self.train_set, batch_size=self.batch_size, num_workers=self.num_workers)

    def val_dataloader(self):
        return GeomDataLoader(self.val_set, batch_size=self.batch_size, num_workers=self.num_workers)

    def test_dataloader(self):
        return GeomDataLoader(self.test_set, batch_size=self.batch_size, num_workers=self.num_workers)<jupyter_output><empty_output><jupyter_text>Networks<jupyter_code>from torch_geometric.nn.models.dimenet import DimeNet as GeoDimeNet

# wraper for error in initialization, for proper initialization torch.no_grad() required
class DimeNet(GeoDimeNet):
    def __init__(self, *args, **kwargs):
        with torch.no_grad():
            super().__init__(*args, **kwargs)
import torch.nn.functional as F
from shared import Module

class DimeNetModule(Module):
    def __init__(self, **kwargs):
        super().__init__()
        self.net = DimeNet(**kwargs)
        
    def step(self, batch, batch_idx):
        z, pos, y, g_batch = batch.z, batch.pos, batch.y, batch.batch
        output = self.net(z, pos, g_batch)
        loss = F.l1_loss(output, y)
        return loss, output.detach(), y.detach()<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint

if __name__ ==  "__main__":
    logger = pl.loggers.tensorboard.TensorBoardLogger("./lightning_logs/", name='dimenet', version='with_gradient_clip')
    early_stop_callback = EarlyStopping(monitor="loss/val", patience=10, verbose=True, mode="min")
    checkpoint_callback = ModelCheckpoint(monitor='loss/val', filename='dimenet-epoch{epoch:02d}-val_loss{val/loss:.2f}',
                                          auto_insert_metric_name=False, save_top_k=10)


    trainer = pl.Trainer(gpus=[0], logger=logger, max_epochs=100,  
                         callbacks=[checkpoint_callback, early_stop_callback],  gradient_clip_val=0.5)
    datamodule = GraphQM9Datamodule()
    model = DimeNetModule(hidden_channels=128, out_channels=12, num_blocks=6,
                            num_bilinear=8, num_spherical=7, num_radial=6,
                            cutoff=5.0, envelope_exponent=5, num_before_skip=1,
                            num_after_skip=2, num_output_layers=3)
if __name__ ==  "__main__":
    trainer.fit(model, datamodule)
    trainer.test(model, datamodule)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10327272727272727, 'fraction_numerical': 0.012363636363636363, 'mean_word_length': 4.132462686567164, 'pattern_counts': {'":': 2, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '141672', 'license': 'MIT', 'n_tokens_mistral': 935, 'n_tokens_neox': 911, 'n_words': 150, 'path': 'train_dimenet.ipynb', 'repo_name': 'a-zapala/quantum_chemistry', 'size': '2750'}