Boilerplate for TorchQuantum Driven Deep Learning Research with PyTorch Lightning and Lightning Fabric
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

from setuptools import setup

rootdir = Path(__file__).parent
long_description = (rootdir / "README.md").read_text()

setup(
    name="quantumlab",
    version="0.0.1",
    package_dir={"": "src"},
    packages=["quantumlab"],
    description="",
    long_description=long_description,
    long_description_content_type="text/markdown",
    author="",
    license="Apache 2.0",
    author_email="",
    classifiers=[
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Topic :: Scientific/Engineering :: Information Analysis",
        "Operating System :: OS Independent",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
    ],
)

setup.py
JustinGoheen/quantum-lab
import os
from pathlib import Path

import typer

FILEPATH = Path(__file__)
PROJECTPATH = FILEPATH.parents[2]
PKGPATH = FILEPATH.parents[1]

app = typer.Typer()
docs_app = typer.Typer()
app.add_typer(docs_app, name="docs")


@app.callback()
def callback() -> None:
    pass


@docs_app.command("build")
def build_docs() -> None:
    import shutil

    os.system("mkdocs build")
    shutil.copyfile(src="README.md", dst="docs/index.md")


@docs_app.command("serve")
def serve_docs() -> None:
    os.system("mkdocs serve")

src/quantumlab/cli/lab.py
JustinGoheen/quantum-lab
import os
from pathlib import Path

# SET PATHS
filepath = Path(__file__)
PROJECTPATH = filepath.parents[1]
LOGSPATH = os.path.join(PROJECTPATH, "logs")
TORCHPROFILERPATH = os.path.join(LOGSPATH, "torch_profiler")
SIMPLEPROFILERPATH = os.path.join(LOGSPATH, "simple_profiler")
CHKPTSPATH = os.path.join(PROJECTPATH, "models", "trials")
MODELPATH = os.path.join(PROJECTPATH, "models", "production", "model.onnx")
PREDSPATH = os.path.join(PROJECTPATH, "data", "predictions", "predictions.pt")
SPLITSPATH = os.path.join(PROJECTPATH, "data", "training_split")
WANDBPATH = os.path.join(PROJECTPATH, "logs", "wandb_logs")
OPTUNAPATH = os.path.join(PROJECTPATH, "logs", "optuna")

src/quantumlab/config.py
JustinGoheen/quantum-lab
from quantumlab.core.module import LabModule
from quantumlab.core.trainer import LabTrainer

src/quantumlab/core/__init__.py
JustinGoheen/quantum-lab
import pytorch_lightning as pl
import torch.nn.functional as F  # noqa: F401
import torchmetrics  # noqa: F401
from torch import optim  # noqa: F401


class LabModule(pl.LightningModule):
    """a custom PyTorch Lightning LightningModule"""

    def __init__(self):
        super().__init__()

    def forward(self, x):
        pass

    def training_step(self, batch):
        pass

    def test_step(self, batch, *args):
        pass

    def validation_step(self, batch, *args):
        pass

    def predict_step(self, batch, batch_idx, dataloader_idx=0):
        pass

    def configure_optimizers(self):
        pass

src/quantumlab/core/module.py
JustinGoheen/quantum-lab
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import pytorch_lightning as pl
import torch
from lightning.pytorch import seed_everything
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.loggers import Logger, TensorBoardLogger
from lightning.pytorch.profilers import Profiler, PyTorchProfiler

from quantumlab import config


class LabTrainer(pl.Trainer):
    def __init__(
        self,
        logger: Optional[Logger] = None,
        profiler: Optional[Profiler] = None,
        callbacks: Optional[List] = [],
        plugins: Optional[List] = [],
        set_seed: bool = True,
        **trainer_init_kwargs: Dict[str, Any]
    ) -> None:
        # SET SEED
        if set_seed:
            seed_everything(config.GLOBALSEED, workers=True)
        super().__init__(
            logger=logger or TensorBoardLogger(config.LOGSPATH, name="tensorboard"),
            profiler=profiler or PyTorchProfiler(dirpath=config.TORCHPROFILERPATH, filename="profiler"),
            callbacks=callbacks + [ModelCheckpoint(dirpath=config.CHKPTSPATH, filename="model")],
            plugins=plugins,
            **trainer_init_kwargs
        )

    def persist_predictions(self, predictions_dir: Optional[Union[str, Path]] = config.PREDSPATH) -> None:
        self.test(ckpt_path="best", datamodule=self.datamodule)
        predictions = self.predict(self.model, self.datamodule.val_dataloader())
        torch.save(predictions, predictions_dir)

src/quantumlab/core/trainer.py
JustinGoheen/quantum-lab
from quantumlab.pipeline.datamodule import LabDataModule

src/quantumlab/pipeline/__init__.py
JustinGoheen/quantum-lab
import multiprocessing
import os
from pathlib import Path

from lightning.pytorch import LightningDataModule
from torch.utils.data import DataLoader, Dataset, random_split

from quantumlab.pipeline.dataset import LabDataset

filepath = Path(__file__)
PROJECTPATH = os.getcwd()
NUMWORKERS = int(multiprocessing.cpu_count() // 2)


class LabDataModule(LightningDataModule):
    def __init__(
        self,
        dataset: Dataset = LabDataset,
        data_dir: str = "data",
        split: bool = True,
        train_size: float = 0.8,
        num_workers: int = NUMWORKERS,
        transforms=None,
    ):
        super().__init__()
        self.data_dir = os.path.join(PROJECTPATH, data_dir, "cache")
        self.dataset = dataset
        self.split = split
        self.train_size = train_size
        self.num_workers = num_workers
        self.transforms = transforms

    def prepare_data(self):
        self.dataset(self.data_dir, download=True)

    def setup(self, stage=None):
        if stage == "fit" or stage is None:
            full_dataset = self.dataset(self.data_dir, train=True, transform=self.transforms)
            train_size = int(len(full_dataset) * self.train_size)
            test_size = len(full_dataset) - train_size
            self.train_data, self.val_data = random_split(full_dataset, lengths=[train_size, test_size])
        if stage == "test" or stage is None:
            self.test_data = self.dataset(self.data_dir, train=False, transform=self.transforms)

    # def teardown(self):
    #     pass

    def train_dataloader(self):
        return DataLoader(self.train_data, num_workers=self.num_workers)

    def test_dataloader(self):
        return DataLoader(self.test_data, num_workers=self.num_workers)

    def val_dataloader(self):
        return DataLoader(self.val_data, num_workers=self.num_workers)

src/quantumlab/pipeline/datamodule.py
JustinGoheen/quantum-lab
class LabDataset:
    """
    Note:
        see below for a basic example of a custom torch dataset
    ```python
    import pandas as pd
    import torch
    from torch.utils.data import Dataset


    class LabDataset(Dataset):
        def __init__(self, features_path, labels_path):
            self.features = pd.read_csv(features_path)
            self.labels = pd.read_csv(labels_path)

        def __len__(self):
            return len(self.labels)

        def __getitem__(self, idx):
            x, y = self.features.iloc[idx], self.labels.iloc[idx]
            return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.float32)
    ```
    """

src/quantumlab/pipeline/dataset.py
JustinGoheen/quantum-lab
from pathlib import Path

import streamlit as st

# create the title
file = Path(__file__)
rootdir = file.parents[1]
long_description = (rootdir / "README.md").read_text()
st.markdown(body=long_description)

streamlit/app.py
JustinGoheen/quantum-lab
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from pathlib import Path

import torch

from quantumlab.pipeline.datamodule import LabDataModule


def test_module_not_abstract():
    _ = LabDataModule()


def test_prepare_data():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    networkpath = Path(__file__).parent
    projectpath = networkpath.parents[0]
    datapath = os.path.join(projectpath, "data", "cache")
    assert "LabDataset" in os.listdir(datapath)


def test_setup():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    datamodule.setup()
    data_keys = ["train_data", "test_data", "val_data"]
    assert all(key in dir(datamodule) for key in data_keys)


def test_trainloader():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    datamodule.setup()
    loader = datamodule.train_dataloader()
    sample = loader.dataset[0][0]
    assert isinstance(sample, torch.Tensor)


def test_testloader():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    datamodule.setup()
    loader = datamodule.test_dataloader()
    sample = loader.dataset[0][0]
    assert isinstance(sample, torch.Tensor)


def test_valloader():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    datamodule.setup()
    loader = datamodule.val_dataloader()
    sample = loader.dataset[0][0]
    assert isinstance(sample, torch.Tensor)

tests/test_datamodule.py
JustinGoheen/quantum-lab
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from quantumlab.core.module import LabModule


def test_module_not_abstract():
    _ = LabModule()


def test_module_forward():
    input_sample = torch.randn((1, 784))
    model = LabModule()
    preds, label = model.forward(input_sample)
    assert preds.shape == input_sample.shape


def test_module_training_step():
    input_sample = torch.randn((1, 784)), 1
    model = LabModule()
    loss = model.training_step(input_sample)
    assert isinstance(loss, torch.Tensor)


def test_optimizer():
    model = LabModule()
    optimizer = model.configure_optimizers()
    optimizer_base_class = optimizer.__class__.__base__.__name__
    assert optimizer_base_class == "Optimizer"

tests/test_network_module.py
JustinGoheen/quantum-lab