From f036fc3c9578765ca0c77263137f133dbb579a4a Mon Sep 17 00:00:00 2001 From: Ishan Upadhyay Date: Sun, 8 Oct 2023 13:24:41 +0530 Subject: [PATCH] Added 'cuda' functionality as well as streamlined the dataloading process for segmentation tasks. --- UniTrain.egg-info/PKG-INFO | 23 ----------------------- UniTrain.egg-info/SOURCES.txt | 10 ---------- UniTrain.egg-info/dependency_links.txt | 1 - UniTrain.egg-info/requires.txt | 5 ----- UniTrain.egg-info/top_level.txt | 1 - UniTrain/utils/classification.py | 14 ++++++++++---- UniTrain/utils/segmentation.py | 26 +++++++++++++++----------- setup.py | 2 +- 8 files changed, 26 insertions(+), 56 deletions(-) delete mode 100644 UniTrain.egg-info/PKG-INFO delete mode 100644 UniTrain.egg-info/SOURCES.txt delete mode 100644 UniTrain.egg-info/dependency_links.txt delete mode 100644 UniTrain.egg-info/requires.txt delete mode 100644 UniTrain.egg-info/top_level.txt diff --git a/UniTrain.egg-info/PKG-INFO b/UniTrain.egg-info/PKG-INFO deleted file mode 100644 index 3b68755..0000000 --- a/UniTrain.egg-info/PKG-INFO +++ /dev/null @@ -1,23 +0,0 @@ -Metadata-Version: 2.1 -Name: UniTrain -Version: 0.1 -Summary: A generalized training framework for Deep Learning Tasks -Home-page: https://github.com/ishan121028/UniTrain -Download-URL: https://github.com/ishan121028/UniTrain/archive/v0.1.tar.gz -Author: Ishan Upadhyay -Author-email: ishan.upadhyay.iitb@gmail.com -Keywords: Deep Learning,Machine Learning,Training Framework -Classifier: Development Status :: 3 - Alpha -Classifier: Intended Audience :: Developers -Classifier: Intended Audience :: Education -Classifier: License :: OSI Approved :: MIT License -Classifier: Programming Language :: Python :: 3 -Classifier: Programming Language :: Python :: 3.7 -Classifier: Programming Language :: Python :: 3.8 -Classifier: Programming Language :: Python :: 3.9 -Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence -Description-Content-Type: text/markdown -License-File: LICENSE - -# UniTrain -UniTrain is an open-source, unified platform for effortless machine learning model training, evaluation, and deployment across diverse tasks. Experience seamless experimentation and model deployment with UniTrain. diff --git a/UniTrain.egg-info/SOURCES.txt b/UniTrain.egg-info/SOURCES.txt deleted file mode 100644 index c1909f2..0000000 --- a/UniTrain.egg-info/SOURCES.txt +++ /dev/null @@ -1,10 +0,0 @@ -LICENSE -README.md -setup.py -UniTrain/__init__.py -UniTrain/train.py -UniTrain.egg-info/PKG-INFO -UniTrain.egg-info/SOURCES.txt -UniTrain.egg-info/dependency_links.txt -UniTrain.egg-info/requires.txt -UniTrain.egg-info/top_level.txt \ No newline at end of file diff --git a/UniTrain.egg-info/dependency_links.txt b/UniTrain.egg-info/dependency_links.txt deleted file mode 100644 index 8b13789..0000000 --- a/UniTrain.egg-info/dependency_links.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/UniTrain.egg-info/requires.txt b/UniTrain.egg-info/requires.txt deleted file mode 100644 index 22d90ea..0000000 --- a/UniTrain.egg-info/requires.txt +++ /dev/null @@ -1,5 +0,0 @@ -torch -torchvision -numpy -pandas -torchsummary diff --git a/UniTrain.egg-info/top_level.txt b/UniTrain.egg-info/top_level.txt deleted file mode 100644 index 571b35c..0000000 --- a/UniTrain.egg-info/top_level.txt +++ /dev/null @@ -1 +0,0 @@ -UniTrain diff --git a/UniTrain/utils/classification.py b/UniTrain/utils/classification.py index bd4744c..54ad233 100644 --- a/UniTrain/utils/classification.py +++ b/UniTrain/utils/classification.py @@ -41,10 +41,8 @@ def get_data_loader(data_dir, batch_size, shuffle=True, transform = None, split= # Create a custom dataset dataset = ClassificationDataset(data_dir, transform=transform) - print(dataset.__len__()) # Create a data loader - print(batch_size) data_loader = DataLoader( dataset, batch_size=batch_size, @@ -86,7 +84,7 @@ def parse_folder(dataset_path): print("An error occurred:", str(e)) return None -def train_model(model, train_data_loader, test_data_loader, num_epochs, learning_rate=0.001, checkpoint_dir='checkpoints', logger=None, device='cpu'): +def train_model(model, train_data_loader, test_data_loader, num_epochs, learning_rate=0.001, checkpoint_dir='checkpoints', logger=None, device=torch.device('cpu')): '''Train a PyTorch model for a classification task. Args: model (nn.Module): Torch model to train. @@ -96,11 +94,17 @@ def train_model(model, train_data_loader, test_data_loader, num_epochs, learning learning_rate (float): Learning rate for the optimizer. checkpoint_dir (str): Directory to save model checkpoints. logger (Logger): Logger to log training details. + device (torch.device): Device to run training on (GPU or CPU). Returns: None ''' + if logger: + logging.basicConfig(level=logging.INFO, format='%(asctime)s - Epoch %(epoch)d - Train Acc: %(train_acc).4f - Val Acc: %(val_acc).4f - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', filename=logger, filemode='w') + logger = logging.getLogger(__name__) + # Define loss function and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) @@ -115,6 +119,9 @@ def train_model(model, train_data_loader, test_data_loader, num_epochs, learning for batch_idx, (inputs, labels) in enumerate(train_data_loader): optimizer.zero_grad() # Zero the parameter gradients + inputs = inputs.to(device) + labels = labels.to(device) + # Forward pass outputs = model(inputs) loss = criterion(outputs, labels) @@ -138,7 +145,6 @@ def train_model(model, train_data_loader, test_data_loader, num_epochs, learning logger.info(f'Epoch {epoch + 1}, Validation Accuracy: {accuracy:.2f}%') - if accuracy > best_accuracy: best_accuracy = accuracy checkpoint_path = os.path.join(checkpoint_dir, f'model_epoch_{epoch + 1}.pth') diff --git a/UniTrain/utils/segmentation.py b/UniTrain/utils/segmentation.py index a63ea4d..d1813fb 100644 --- a/UniTrain/utils/segmentation.py +++ b/UniTrain/utils/segmentation.py @@ -7,9 +7,11 @@ from torchvision import transforms from ..dataset.segmentation import SegmentationDataset import torchsummary +import logging +import glob -def get_data_loader(image_paths:list, mask_paths:list, batch_size:int, shuffle:bool=True, transform=None) -> DataLoader: - """ +def get_data_loader(data_dir: str, batch_size:int, shuffle:bool=True, transform=None, split='train') -> DataLoader: + """, Create and return a data loader for a custom dataset. Args: @@ -28,6 +30,9 @@ def get_data_loader(image_paths:list, mask_paths:list, batch_size:int, shuffle:b transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) # Normalize with ImageNet stats ]) + image_paths = glob.glob(os.path.join(data_dir, split, 'images', '*')) + mask_paths = glob.glob(os.path.join(data_dir, split, 'masks', '*')) + # Create a custom dataset dataset = SegmentationDataset(image_paths=image_paths, mask_paths=mask_paths, transform=transform) @@ -83,7 +88,7 @@ def parse_folder(dataset_path): return False -def train_unet(model, train_data_loader, test_data_loader, num_epochs, learning_rate, checkpoint_dir, logger=None, iou=False, device='cpu'): +def train_unet(model, train_data_loader, test_data_loader, num_epochs, learning_rate, checkpoint_dir, logger=None, iou=False, device=torch.device('cpu')) -> None: '''Train the model using the given train and test data loaders. Args: @@ -95,17 +100,17 @@ def train_unet(model, train_data_loader, test_data_loader, num_epochs, learning_ checkpoint_dir (str): Directory to save model checkpoints. logger (Logger): Logger to log training information. iou (bool): Whether to calculate the IOU score. + device (torch.device): Device to run training on (GPU or CPU). Returns: None ''' - if device == 'cpu': - device = torch.device('cpu') - elif device == 'cuda': - device = torch.device('cuda') - else: - print(f"{device} is not a valid device.") - return None + + if logger: + logging.basicConfig(level=logging.INFO, format='%(asctime)s - Epoch %(epoch)d - Train Acc: %(train_acc).4f - Val Acc: %(val_acc).4f - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', filename=logger, filemode='w') + logger = logging.getLogger(__name__) + criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) @@ -182,7 +187,6 @@ def iou_score(output, target): Returns: float: The average IoU score. ''' - smooth = 1e-6 output = output.argmax(1) intersection = (output & target).float().sum((1, 2)) diff --git a/setup.py b/setup.py index 9bd5acf..021b1b6 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ setup( name='UniTrain', - version='0.2.2', + version='0.2.3', author='Ishan Upadhyay', author_email='ishan.upadhyay.iitb@gmail.com', description='A generalized training framework for Deep Learning Tasks',