blog_8.py

def instantiate_datamodule(args):
    datamodule = CoraDataset(
        num_workers=args.num_workers,
        batch_size=args.batch_size,
        drop_last=args.drop_last,
        pin_memory=args.pin_memory,
        num_layers=args.num_layers,
    )
    return datamodule

def instantiate_model(args, datamodule):
    model = DNAConvNet(
        num_layers=args.num_layers,
        hidden_channels=args.hidden_channels,
        heads=args.heads,
        groups=args.groups,
        dropout=args.dropout,
        # provide dataset specific arguments
        **datamodule.hyper_parameters,
    )
    # Attached datamodule function to model
    model.gather_data = datamodule.gather_data
    return model

def run(args):
    datamodule: LightningDataModule = instantiate_datamodule(args)
    model: LightningModule = instantiate_model(args, datamodule)
    trainer = Trainer.from_argparse_args(args)
    trainer.fit(model, datamodule)
    trainer.test()

if __name__ == "__main__":
    parser = ArgumentParser(description="Pytorch Geometric Example")
    parser = Trainer.add_argparse_args(parser)
    parser = CoraDataset.add_argparse_args(parser)
    parser = DNAConvNet.add_argparse_args(parser)
    cmd_line = '--max_epochs 1'.split(' ')
    run(parser.parse_args(cmd_line))