def __init__(self, lr: float = 0.01, num_classes: int = 19, num_layers: int = 5, features_start: int = 64, bilinear: bool = False): """ Basic model for semantic segmentation. Uses UNet architecture by default. The default parameters in this model are for the KITTI dataset. Note, if you'd like to use this model as is, you will first need to download the KITTI dataset yourself. You can download the dataset `here. <http://www.cvlibs.net/datasets/kitti/eval_semseg.php?benchmark=semantics2015>`_ Implemented by: - `Annika Brundyn <https://github.com/annikabrundyn>`_ Args: num_layers: number of layers in each side of U-net (default 5) features_start: number of features in first layer (default 64) bilinear: whether to use bilinear interpolation (True) or transposed convolutions (default) for upsampling. lr: learning (default 0.01) """ super().__init__() self.num_classes = num_classes self.num_layers = num_layers self.features_start = features_start self.bilinear = bilinear self.lr = lr self.net = UNet(num_classes=num_classes, num_layers=self.num_layers, features_start=self.features_start, bilinear=self.bilinear)
class SemSegment(pl.LightningModule): def __init__(self, lr: float = 0.01, num_classes: int = 19, num_layers: int = 5, features_start: int = 64, bilinear: bool = False): """ Basic model for semantic segmentation. Uses UNet architecture by default. The default parameters in this model are for the KITTI dataset. Note, if you'd like to use this model as is, you will first need to download the KITTI dataset yourself. You can download the dataset `here. <http://www.cvlibs.net/datasets/kitti/eval_semseg.php?benchmark=semantics2015>`_ Implemented by: - `Annika Brundyn <https://github.com/annikabrundyn>`_ Args: num_layers: number of layers in each side of U-net (default 5) features_start: number of features in first layer (default 64) bilinear: whether to use bilinear interpolation (True) or transposed convolutions (default) for upsampling. lr: learning (default 0.01) """ super().__init__() self.num_classes = num_classes self.num_layers = num_layers self.features_start = features_start self.bilinear = bilinear self.lr = lr self.net = UNet(num_classes=num_classes, num_layers=self.num_layers, features_start=self.features_start, bilinear=self.bilinear) def forward(self, x): return self.net(x) def training_step(self, batch, batch_nb): img, mask = batch img = img.float() mask = mask.long() out = self(img) loss_val = F.cross_entropy(out, mask, ignore_index=250) log_dict = {'train_loss': loss_val} return {'loss': loss_val, 'log': log_dict, 'progress_bar': log_dict} def validation_step(self, batch, batch_idx): img, mask = batch img = img.float() mask = mask.long() out = self(img) loss_val = F.cross_entropy(out, mask, ignore_index=250) return {'val_loss': loss_val} def validation_epoch_end(self, outputs): loss_val = torch.stack([x['val_loss'] for x in outputs]).mean() log_dict = {'val_loss': loss_val} return { 'log': log_dict, 'val_loss': log_dict['val_loss'], 'progress_bar': log_dict } def configure_optimizers(self): opt = torch.optim.Adam(self.net.parameters(), lr=self.lr) sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10) return [opt], [sch] @staticmethod def add_model_specific_args(parent_parser): parser = ArgumentParser(parents=[parent_parser], add_help=False) parser.add_argument("--lr", type=float, default=0.01, help="adam: learning rate") parser.add_argument("--num_layers", type=int, default=5, help="number of layers on u-net") parser.add_argument("--features_start", type=float, default=64, help="number of features in first layer") parser.add_argument( "--bilinear", action='store_true', default=False, help="whether to use bilinear interpolation or transposed") return parser