From 49c03f7dc26adc8f882977fadac9f1151e01f8bc Mon Sep 17 00:00:00 2001 From: Wei Ji <23487320+weiji14@users.noreply.github.com> Date: Sat, 2 Nov 2024 09:52:22 +1300 Subject: [PATCH] :alien: Update CTCoreUnet to handle newer pytorch* library versions Fix the following errors - AttributeError: type object 'Trainer' has no attribute 'add_argparse_args' (need to change from Pytorch Lightning 1.0 to Lightning 2.0 style, see https://github.com/Lightning-AI/pytorch-lightning/issues/19905) - ValueError: Cannot find 8ef2e2d423b67b53ec8113fc71a9b968bb0f66e7 in https://github.com/mateuszbuda/brain-segmentation-pytorch (change to use v1.0 tag) - AttributeError: module 'torchmetrics' has no attribute 'IoU' (renamed to JaccardIndex, see https://github.com/Lightning-AI/torchmetrics/pull/662) Also gitignoring __pycache__/ folder. --- .gitignore | 5 +++++ README.md | 5 +++-- ctcorenet/ctcoredata.py | 2 +- ctcorenet/ctcoreunet.py | 35 ++++++++++++++++++++++------------- 4 files changed, 31 insertions(+), 16 deletions(-) diff --git a/.gitignore b/.gitignore index 170849b..d7e9cb9 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,8 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +.pytest_cache/ +*.py[cod] + # Raw data tracked by dvc data/** !data/**/ diff --git a/README.md b/README.md index 84c18bc..0d8b7ed 100644 --- a/README.md +++ b/README.md @@ -72,9 +72,9 @@ This will load the image data stored in `data/train`, perform the training (minimize loss between img.png and label.png), and produce some outputs. More advanced users can customize the training, e.g. to be more deterministic, -running for only x epochs, train on a GPU using 16-bit precision, etc, like so: +running for only x epochs, train on an CUDA GPU using 16-bit precision, etc, like so: - python ctcorenet/ctcoreunet.py --deterministic=True --max_epochs=3 --gpus=1 --precision=16 + python ctcorenet/ctcoreunet.py --deterministic=True --max_epochs=3 --accelerator=gpu --devices=1 --precision=16 More options to customize the training can be found by running `python ctcorenet/ctcoreunet.py --help`. @@ -86,6 +86,7 @@ data version control ([DVC](https://github.com/iterative/dvc/)) library which stores all the commands and input/intermediate/output data assets used. This makes it easy to reproduce the entire pipeline using a single command + dvc pull dvc repro This command will perform all the data preparation and model training steps. diff --git a/ctcorenet/ctcoredata.py b/ctcorenet/ctcoredata.py index e18aedb..faed84a 100644 --- a/ctcorenet/ctcoredata.py +++ b/ctcorenet/ctcoredata.py @@ -8,7 +8,7 @@ import os import typing -import pytorch_lightning as pl +import lightning.pytorch as pl import torch import torchvision from torch.nn import functional as F diff --git a/ctcorenet/ctcoreunet.py b/ctcorenet/ctcoreunet.py index e1e5764..fa4a0f0 100644 --- a/ctcorenet/ctcoreunet.py +++ b/ctcorenet/ctcoreunet.py @@ -6,14 +6,12 @@ """ import argparse -import typing -import pytorch_lightning as pl +import lightning.pytorch as pl import torch import torchmetrics import torchvision from ctcoredata import CTCoreDataModule -from torch.nn import functional as F torch.use_deterministic_algorithms(mode=True) @@ -39,18 +37,19 @@ def __init__(self): # Unet model # https://github.com/mateuszbuda/brain-segmentation-pytorch#pytorch-hub self.unet = torch.hub.load( - repo_or_dir="mateuszbuda/brain-segmentation-pytorch:8ef2e2d423b67b53ec8113fc71a9b968bb0f66e7", + repo_or_dir="mateuszbuda/brain-segmentation-pytorch:v1.0", model="unet", in_channels=1, out_channels=1, # binary classification init_features=32, pretrained=False, + trust_repo="check", ) # self.output_conv = torch.nn.Conv2d(32, 1, 1, 1) # Intersection over Union (IoU) metric - self.iou = torchmetrics.IoU(num_classes=2) + self.iou = torchmetrics.JaccardIndex(task="binary", num_classes=2) def forward(self, x) -> torch.Tensor: """ @@ -135,7 +134,7 @@ def configure_optimizers(self): return torch.optim.Adam(params=self.parameters(), lr=0.01) -def cli_main(): +def cli_main(hparams): """ Command line interface to run the CTCoreNet model. Based on https://github.com/PyTorchLightning/deep-learning-project-template @@ -161,11 +160,6 @@ def cli_main(): More options can be found by using `python ctcorenet/ctcoreunet.py --help`. Happy training! """ - ## Parse arguments from command-line - parser = argparse.ArgumentParser() - parser = pl.Trainer.add_argparse_args(parent_parser=parser) - args = parser.parse_args() - # Set a seed to control for randomness pl.seed_everything(42) @@ -176,7 +170,13 @@ def cli_main(): model = CTCoreNet() # Training - trainer = pl.Trainer.from_argparse_args(args=args) + trainer = pl.Trainer( + accelerator=hparams.accelerator, + deterministic=hparams.deterministic, + max_epochs=hparams.max_epochs, + devices=hparams.devices, + precision=hparams.precision, + ) trainer.fit(model=model, datamodule=ctcoredatamodule) # Export Model to ONNX format @@ -189,4 +189,13 @@ def cli_main(): if __name__ == "__main__": - cli_main() + ## Parse arguments from command-line + parser = argparse.ArgumentParser() + parser.add_argument("--accelerator", default="auto") # "gpu" or "mps" (mac M1 chip) + parser.add_argument("--deterministic", default=None) + parser.add_argument("--max_epochs", default=3, type=int) + parser.add_argument("--devices", default="auto") + parser.add_argument("--precision", default=None) + args = parser.parse_args() + + cli_main(args)