<gh_stars>0
import pytorch_lightning as pl
import torch
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader

from config.custom_config import custom_parser_config
from model.MaskRCNN import get_model_instance_segmentation
from pidata.pidata import pi_parser
from piutils.piutils import pi_log
from utils.checkpoint_utils import PeriodicCheckpoint

logger = pi_log.get_logger(__name__)
logger.propagate = False

# Setting seed for reproducibility
seed = 666
torch.manual_seed(seed)

Tensor = torch.tensor
Module = torch.nn.Module


class SegmentationModule(pl.LightningModule):
    """
    Pytorch Lightning module for training
    different segmentation models.
    """

    def __init__(self, config_data: dict, train_mode: bool = True,
                 batch_size: int = 5, gpu: int = 1,
                 epochs: int = 50, lr: float = 0.003) -> None:
        """
        :param config_data: Config for data, utils.
        :param train_mode: Sets model to training or run_inference mode.
        :param batch_size: Batch size during training.
        :param epochs: Total epochs for training. (default 50)
        :param lr: Learning rate (default 0.003)
        :param gpu: Set total gpus to use (default 1)
        """
        super(SegmentationModule, self).__init__()

        # (1)-> Testing custom model
        # self.model = UNET(3, 4)

        # (2)-> Default Mask-RCNN model
        self.model = get_model_instance_segmentation(4)

        # (3)-> Model from 'segmentation_models_pytorch' library
        # self.model = smp.Unet(
        #     encoder_name="mobilenet_v2",  # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
        #     encoder_weights="imagenet",  # use `imagenet` pre-trained weights for encoder initialization
        #     in_channels=3,  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
        #     classes=4,  # model output channels (number of classes in your dataset)
        # )

        self.model.train(train_mode)

        self.val_loader, self.train_loader = None, None
        self.num_train_imgs, self.num_val_imgs = None, None
        self.trainer, self.curr_device = None, None

        # Model checkpoint saving every 10 steps
        self.periodic_chkp = PeriodicCheckpoint(10)
        self.save_hyperparameters()

        self.config_data = config_data
        self.loss_fn = CrossEntropyLoss()
        self.batch_size = batch_size
        self.epochs = epochs
        self.samples = 1000
        self.learning_rate = lr
        self.gpu = gpu

    def forward(self, input: Tensor, **kwargs) -> Tensor:
        return self.model(input)

    def training_step(self, batch, batch_idx):
        print("Batch ==>", batch_idx, len(batch))

        inputs, labels = batch

        loss_dict = self.model(inputs, labels)
        print("Training Loss :: ", type(loss_dict), len(loss_dict))
        losses = sum(loss for loss in loss_dict.values())

        return losses

    def validation_step(self, batch, batch_idx):
        print("Batch ==>", batch_idx, len(batch))

        inputs, labels = batch

        loss_dict = self.model(inputs, labels)
        print("Loss :: ", type(loss_dict), len(loss_dict))

        return loss_dict

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
        return optimizer

    def train_dataloader(self):
        training_data_parser = pi_parser.PiParser(
            config=custom_parser_config,
            split_name="train",
            num_samples=self.samples,  # number of samples to be drawn, set to a multiple of the batch size
            numpy_to_tensor_func=torch.from_numpy,
            # framework-dependent, e.g. torch.from_numpy (PyTorch), if None, the returned type is numpy.ndarray
        )

        self.train_loader = DataLoader(training_data_parser,
                                       batch_size=self.batch_size,
                                       shuffle=True,
                                       num_workers=4,
                                       collate_fn=self.collate_fn  # custom collate
                                       )
        self.num_train_imgs = len(self.train_loader)
        return self.train_loader

    def val_dataloader(self):
        val_data_parser = pi_parser.PiParser(
            config=custom_parser_config,
            split_name="val",
            num_samples=self.samples,  # number of samples to be drawn, set to a multiple of the batch size
            numpy_to_tensor_func=torch.from_numpy,
            # framework-dependent, e.g. torch.from_numpy (PyTorch), if None, the returned type is numpy.ndarray
        )

        self.val_loader = DataLoader(val_data_parser,
                                     batch_size=self.batch_size,
                                     shuffle=True,
                                     num_workers=4,
                                     collate_fn=self.collate_fn  # custom collate
                                     )
        self.num_val_imgs = len(self.val_loader)
        return self.val_loader

    def train_model(self):
        self.trainer = pl.Trainer(gpus=self.gpu,
                                  max_epochs=self.epochs,
                                  callbacks=self.periodic_chkp,
                                  weights_summary="full",
                                  # overfit_batches=2,
                                  # accumulate_grad_batches=2,
                                  check_val_every_n_epoch=10,
                                  stochastic_weight_avg=True
                                  )
        self.trainer.fit(self,
                         self.train_dataloader(),
                         self.val_dataloader())

    def collate_fn(self, batch):
        image_list, target_list = [], []
        for item in batch:
            images, targets = item

            image_list.append(images)
            target_list.append(targets)

        # print("Img ->", len(image_list), type(image_list))
        # print("Targets ->", len(target_list), type(target_list))
        return image_list, target_list


if __name__ == "__main__":
    model_trainer = SegmentationModule(config_data=custom_parser_config,
                                       train_mode=True,
                                       lr=0.0001,
                                       batch_size=2,
                                       epochs=10,
                                       gpu=1)
    model_trainer.train_model()

import pytorch_lightning as pl
import torch
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader

from config.custom_config import custom_parser_config
from model.MaskRCNN import get_model_instance_segmentation
from pidata.pidata import pi_parser
from piutils.piutils import pi_log
from utils.checkpoint_utils import PeriodicCheckpoint

logger = pi_log.get_logger(__name__)
logger.propagate = False

# Setting seed for reproducibility
seed = 666
torch.manual_seed(seed)

Tensor = torch.tensor
Module = torch.nn.Module


class SegmentationModule(pl.LightningModule):
    """
    Pytorch Lightning module for training
    different segmentation models.
    """

    def __init__(self, config_data: dict, train_mode: bool = True,
                 batch_size: int = 5, gpu: int = 1,
                 epochs: int = 50, lr: float = 0.003) -> None:
        """
        :param config_data: Config for data, utils.
        :param train_mode: Sets model to training or run_inference mode.
        :param batch_size: Batch size during training.
        :param epochs: Total epochs for training. (default 50)
        :param lr: Learning rate (default 0.003)
        :param gpu: Set total gpus to use (default 1)
        """
        super(SegmentationModule, self).__init__()

        # (1)-> Testing custom model
        # self.model = UNET(3, 4)

        # (2)-> Default Mask-RCNN model
        self.model = get_model_instance_segmentation(4)

        # (3)-> Model from 'segmentation_models_pytorch' library
        # self.model = smp.Unet(
        #     encoder_name="mobilenet_v2",  # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
        #     encoder_weights="imagenet",  # use `imagenet` pre-trained weights for encoder initialization
        #     in_channels=3,  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
        #     classes=4,  # model output channels (number of classes in your dataset)
        # )

        self.model.train(train_mode)

        self.val_loader, self.train_loader = None, None
        self.num_train_imgs, self.num_val_imgs = None, None
        self.trainer, self.curr_device = None, None

        # Model checkpoint saving every 10 steps
        self.periodic_chkp = PeriodicCheckpoint(10)
        self.save_hyperparameters()

        self.config_data = config_data
        self.loss_fn = CrossEntropyLoss()
        self.batch_size = batch_size
        self.epochs = epochs
        self.samples = 1000
        self.learning_rate = lr
        self.gpu = gpu

    def forward(self, input: Tensor, **kwargs) -> Tensor:
        return self.model(input)

    def training_step(self, batch, batch_idx):
        print("Batch ==>", batch_idx, len(batch))

        inputs, labels = batch

        loss_dict = self.model(inputs, labels)
        print("Training Loss :: ", type(loss_dict), len(loss_dict))
        losses = sum(loss for loss in loss_dict.values())

        return losses

    def validation_step(self, batch, batch_idx):
        print("Batch ==>", batch_idx, len(batch))

        inputs, labels = batch

        loss_dict = self.model(inputs, labels)
        print("Loss :: ", type(loss_dict), len(loss_dict))

        return loss_dict

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
        return optimizer

    def train_dataloader(self):
        training_data_parser = pi_parser.PiParser(
            config=custom_parser_config,
            split_name="train",
            num_samples=self.samples,  # number of samples to be drawn, set to a multiple of the batch size
            numpy_to_tensor_func=torch.from_numpy,
            # framework-dependent, e.g. torch.from_numpy (PyTorch), if None, the returned type is numpy.ndarray
        )

        self.train_loader = DataLoader(training_data_parser,
                                       batch_size=self.batch_size,
                                       shuffle=True,
                                       num_workers=4,
                                       collate_fn=self.collate_fn  # custom collate
                                       )
        self.num_train_imgs = len(self.train_loader)
        return self.train_loader

    def val_dataloader(self):
        val_data_parser = pi_parser.PiParser(
            config=custom_parser_config,
            split_name="val",
            num_samples=self.samples,  # number of samples to be drawn, set to a multiple of the batch size
            numpy_to_tensor_func=torch.from_numpy,
            # framework-dependent, e.g. torch.from_numpy (PyTorch), if None, the returned type is numpy.ndarray
        )

        self.val_loader = DataLoader(val_data_parser,
                                     batch_size=self.batch_size,
                                     shuffle=True,
                                     num_workers=4,
                                     collate_fn=self.collate_fn  # custom collate
                                     )
        self.num_val_imgs = len(self.val_loader)
        return self.val_loader

    def train_model(self):
        self.trainer = pl.Trainer(gpus=self.gpu,
                                  max_epochs=self.epochs,
                                  callbacks=self.periodic_chkp,
                                  weights_summary="full",
                                  # overfit_batches=2,
                                  # accumulate_grad_batches=2,
                                  check_val_every_n_epoch=10,
                                  stochastic_weight_avg=True
                                  )
        self.trainer.fit(self,
                         self.train_dataloader(),
                         self.val_dataloader())

    def collate_fn(self, batch):
        image_list, target_list = [], []
        for item in batch:
            images, targets = item

            image_list.append(images)
            target_list.append(targets)

        # print("Img ->", len(image_list), type(image_list))
        # print("Targets ->", len(target_list), type(target_list))
        return image_list, target_list


if __name__ == "__main__":
    model_trainer = SegmentationModule(config_data=custom_parser_config,
                                       train_mode=True,
                                       lr=0.0001,
                                       batch_size=2,
                                       epochs=10,
                                       gpu=1)
    model_trainer.train_model()
# Setting seed for reproducibility Pytorch Lightning module for training     different segmentation models. :param config_data: Config for data, utils.         :param train_mode: Sets model to training or run_inference mode.         :param batch_size: Batch size during training.         :param epochs: Total epochs for training. (default 50)         :param lr: Learning rate (default 0.003)         :param gpu: Set total gpus to use (default 1) # (1)-> Testing custom model # self.model = UNET(3, 4) # (2)-> Default Mask-RCNN model # (3)-> Model from 'segmentation_models_pytorch' library # self.model = smp.Unet( #     encoder_name="mobilenet_v2",  # choose encoder, e.g. mobilenet_v2 or efficientnet-b7 #     encoder_weights="imagenet",  # use `imagenet` pre-trained weights for encoder initialization #     in_channels=3,  # model input channels (1 for gray-scale images, 3 for RGB, etc.) #     classes=4,  # model output channels (number of classes in your dataset) # ) # Model checkpoint saving every 10 steps # number of samples to be drawn, set to a multiple of the batch size # framework-dependent, e.g. torch.from_numpy (PyTorch), if None, the returned type is numpy.ndarray # custom collate # number of samples to be drawn, set to a multiple of the batch size # framework-dependent, e.g. torch.from_numpy (PyTorch), if None, the returned type is numpy.ndarray # custom collate # overfit_batches=2, # accumulate_grad_batches=2, # print("Img ->", len(image_list), type(image_list)) # print("Targets ->", len(target_list), type(target_list))