class TestModel(BoringModel): def __init__(self): super().__init__() self.metric = MetricCollection([SumMetric(), DiffMetric()]) self.sum = 0.0 self.diff = 0.0 def training_step(self, batch, batch_idx): x = batch metric_vals = self.metric(x.sum()) self.sum += x.sum() self.diff -= x.sum() self.log_dict({f"{k}_step": v for k, v in metric_vals.items()}) return self.step(x) def training_epoch_end(self, outputs): metric_vals = self.metric.compute() self.log_dict({f"{k}_epoch": v for k, v in metric_vals.items()})
class Ranker(LightningModule): """Base class for rankers. Implements AP, RR and nDCG for validation and testing. This class needs to be extended and the following methods must be implemented: * forward * configure_optimizers (alternatively, this can be implemented in the data module) """ def __init__( self, training_mode: TrainingMode = TrainingMode.POINTWISE, pairwise_loss_margin: float = 1.0, ) -> None: """Constructor. Args: training_mode (TrainingMode, optional): How to train the model. Defaults to TrainingMode.POINTWISE. pairwise_loss_margin (float, optional): Margin used in pairwise loss. Defaults to 1.0. """ super().__init__() self.training_mode = training_mode self.pairwise_loss_margin = pairwise_loss_margin self.bce = torch.nn.BCEWithLogitsLoss() metrics = [RetrievalMAP, RetrievalMRR, RetrievalNormalizedDCG] self.val_metrics = MetricCollection( [M(compute_on_step=False) for M in metrics], prefix="val_", ) self.test_metrics = MetricCollection( [M(compute_on_step=False) for M in metrics], prefix="test_", ) def training_step( self, batch: Union[PointwiseTrainingBatch, PairwiseTrainingBatch], batch_idx: int, ) -> torch.Tensor: """Train a single batch. Args: batch (Union[PointwiseTrainingBatch, PairwiseTrainingBatch]): A training batch. batch_idx (int): Batch index. Returns: torch.Tensor: Training loss. """ if self.training_mode == TrainingMode.POINTWISE: model_batch, labels, _ = batch loss = self.bce(self(model_batch).flatten(), labels.flatten()) elif self.training_mode == TrainingMode.PAIRWISE: pos_model_batch, neg_model_batch, _ = batch pos_outputs = torch.sigmoid(self(pos_model_batch)) neg_outputs = torch.sigmoid(self(neg_model_batch)) loss = torch.mean( torch.clamp(self.pairwise_loss_margin - pos_outputs + neg_outputs, min=0)) self.log("train_loss", loss) return loss def validation_step(self, batch: ValTestBatch, batch_idx: int) -> Dict[str, torch.Tensor]: """Process a validation batch. The returned query IDs are internal IDs. Args: batch (ValTestBatch): A validation batch. batch_idx (int): Batch index. Returns: Dict[str, torch.Tensor]: Query IDs, scores and labels. """ model_batch, q_ids, labels = batch return { "q_ids": q_ids, "scores": self(model_batch).flatten(), "labels": labels } def validation_step_end(self, step_results: Dict[str, torch.Tensor]) -> None: """Update the validation metrics. Args: step_results (Dict[str, torch.Tensor]): Results from a validation step. """ self.val_metrics( step_results["scores"], step_results["labels"], indexes=step_results["q_ids"], ) def validation_epoch_end( self, val_results: Iterable[Dict[str, torch.Tensor]]) -> None: """Compute validation metrics. Args: val_results (Iterable[Dict[str, torch.Tensor]]): Results of the validation steps. """ for metric, value in self.val_metrics.compute().items(): self.log(metric, value, sync_dist=True) self.val_metrics.reset() def test_step(self, batch: ValTestBatch, batch_idx: int) -> Dict[str, torch.Tensor]: """Process a test batch. The returned query IDs are internal IDs. Args: batch (ValTestBatch): A validation batch. batch_idx (int): Batch index. Returns: Dict[str, torch.Tensor]: Query IDs, scores and labels. """ model_batch, q_ids, labels = batch return { "q_ids": q_ids, "scores": self(model_batch).flatten(), "labels": labels } def test_step_end(self, step_results: Dict[str, torch.Tensor]) -> None: """Update the test metrics. Args: step_results (Dict[str, torch.Tensor]): Results from a test step. """ self.test_metrics( step_results["scores"], step_results["labels"], indexes=step_results["q_ids"], ) def test_epoch_end( self, test_results: Iterable[Dict[str, torch.Tensor]]) -> None: """Compute test metrics. Args: test_results (Iterable[Dict[str, torch.Tensor]]): Results of the test steps. """ for metric, value in self.test_metrics.compute().items(): self.log(metric, value, sync_dist=True) self.test_metrics.reset() def predict_step(self, batch: PredictionBatch, batch_idx: int) -> Dict[str, torch.Tensor]: """Compute scores for a prediction batch. Args: batch (PredictionBatch): Inputs. batch_idx (int): Batch index. dataloader_idx (int): DataLoader index. Returns: Dict[str, torch.Tensor]: Indices and scores. """ indices, model_inputs = batch return {"indices": indices, "scores": self(model_inputs).flatten()}
class ImageRegression(BaseModel): """ Model for image regression. This is a configurable class composed by a backbone (see solarnet.models.backbone.py) and a regressor head (actually a classifier with 1 output). It is also a LightningModule and nn.Module. """ def __init__( self, n_channel: int = 1, learning_rate: float = 1e-4, backbone: Union[str, nn.Module] = "simple-cnn", backbone_output_size: int = 0, n_hidden: int = 512, dropout: float = 0.2, loss_fn: str = "mse", lr_scheduler: bool = False, lr_scheduler_warmup_steps: int = 100, lr_scheduler_total_steps: int = 0, **kwargs, ): super().__init__() self.save_hyperparameters() if isinstance(backbone, str): self.backbone, backbone_output_size = get_backbone( backbone, channels=n_channel, dropout=dropout, **kwargs) self.regressor = Classifier(backbone_output_size, 1, n_hidden, dropout) if loss_fn == "mse": self.loss_fn = nn.MSELoss() elif loss_fn == "mae": self.loss_fn = nn.L1Loss() # MAE else: raise RuntimeError("Undefined loss function") self.test_metrics = MetricCollection([ MeanAbsoluteError(), MeanSquaredError(), ]) @property def backbone_name(self) -> str: if isinstance(self.hparams.backbone, str): return self.hparams.backbone else: return type(self.hparams.backbone).__name__ @property def output_size(self) -> int: return 1 @auto_move_data def forward(self, image): return self.regressor(self.backbone(image)) def training_step(self, batch, batch_id): return self.step(batch, step_type="train") def validation_step(self, batch, batch_id): return self.step(batch, step_type="val") def step(self, batch, step_type: str): image, y = batch y_pred = self(image) y_pred = torch.flatten(y_pred) y = y.float() loss = self.loss_fn(y_pred, y) self.log(f"{step_type}_loss", loss, prog_bar=True, sync_dist=True) return loss def test_step(self, batch, batch_idx): image, y = batch y_pred = self(image) y_pred = torch.flatten(y_pred) y = log_min_max_inverse_scale(y) y_pred = log_min_max_inverse_scale(y_pred) self.test_metrics(y_pred, y) def test_epoch_end(self, outs): test_metrics = self.test_metrics.compute() self.log("test_mae", test_metrics["MeanAbsoluteError"]) self.log("test_mse", test_metrics["MeanSquaredError"]) def configure_optimizers(self): logger.info(f"configure_optimizers lr={self.hparams.learning_rate}") optimizer = optim.Adam( filter(lambda p: p.requires_grad, self.parameters()), lr=self.hparams.learning_rate, ) if not self.hparams.lr_scheduler: return optimizer scheduler = optim.lr_scheduler.LambdaLR( optimizer, linear_warmup_decay(self.hparams.lr_scheduler_warmup_steps, self.hparams.lr_scheduler_total_steps, cosine=True), ) return ({ "optimizer": optimizer, "lr_scheduler": { "scheduler": scheduler, "interval": "step", "frequency": 1, }, }, )
def main(csv_file, data_path): gpu = 'cuda:0' if torch.cuda.is_available() else 'cpu' data_set = CSVDataset(csv_file, data_path) tb_logger = TensorBoardLogger(save_dir='./logs/') data_len = len(data_set) train_len, val_len = int(0.6 * data_len), int(0.2 * data_len) test_len = data_len - (train_len + val_len) train_set, val_set, test_set = random_split( data_set, (train_len, val_len, test_len) ) train_loader = DataLoader(train_set, batch_size=16, shuffle=True) val_loader = DataLoader(val_set, batch_size=16) test_loader = DataLoader(test_set, batch_size=16) model = ClassifierBackBone().to(gpu) optimizer = torch.optim.Adam(model.parameters(), lr=0.1) criterion = torch.nn.BCELoss() writer = tensorboard.SummaryWriter('./logs') writer.add_graph(model, input_to_model=torch.randn(1, 3, 400, 400)) train_collection = MetricCollection([Accuracy(compute_on_step=False)]) val_collection = MetricCollection([Accuracy(compute_on_step=False)]) for ep in range(1, 1000): loss_val = 0 with tqdm.tqdm(train_loader, unit="batch") as train_epoch: for idx, (inp, label) in enumerate(train_epoch): train_epoch.set_description(f'Train: {ep}') inp, label = inp.to(gpu), label.to(gpu) optimizer.zero_grad() out = model(inp) loss = criterion(out, label) loss.backward() optimizer.step() loss_val += loss.item() out, label = torch.round(out).to(int).to('cpu'), label.to(int).to('cpu') train_collection(out, label) train_epoch.set_postfix(loss=loss_val / idx) writer.add_scalars('Training', train_collection.compute(), ep) train_collection.reset() val_loss_val = 0 with tqdm.tqdm(val_loader, unit="batch") as val_epoch: for idx, (inp, label) in enumerate(val_epoch): with torch.no_grad(): inp, label = inp.to(gpu), label.to(gpu) out = model(inp) loss = criterion(out, label) val_loss_val += loss.item() out, label = torch.round(out).to(int).to('cpu'), label.to(int).to('cpu') val_collection(out, label) train_epoch.set_postfix(loss=loss_val / idx) writer.add_scalars('Validation', val_collection.compute(), ep) val_collection.reset() writer.add_scalars('Loss', { 'training': loss_val / len(train_loader), 'validation': val_loss_val / len(val_loader) }, ep) break
class ClassifierBackBone(pl.LightningModule): def __init__(self): super(ClassifierBackBone, self).__init__() self.back_bone = nn.Sequential(nn.Conv2d(3, 32, (7, 7), stride=(2, 2)), ResidualBlock(32, 32), ResidualBlock(32, 64), ResidualBottleneck(64, 2), ResidualBlock(64, 64), ResidualBlock(64, 128), ResidualBottleneck(128, 2), ResidualBlock(128, 128), ResidualBlock(128, 256), ResidualBottleneck(256, 2), ResidualBlock(256, 256), nn.Flatten(), nn.Linear(256 * 25 * 25, 1), nn.Sigmoid()) self.criterion = torch.nn.BCELoss() self.train_metrics = MetricCollection({ 'train_accuracy': Accuracy(compute_on_step=False), 'train_precision': Precision(compute_on_step=False), 'train_recall': Recall(compute_on_step=False), }) self.val_metrics = MetricCollection({ 'val_accuracy': Accuracy(compute_on_step=False), 'val_precision': Precision(compute_on_step=False), 'val_recall': Recall(compute_on_step=False) }) def forward(self, x): return self.back_bone(x) def configure_optimizers(self): optimizer_func = torch.optim.Adam(self.parameters(), lr=1e-2) return optimizer_func def training_step(self, train_batch, batch_idx): inp, label = train_batch out = self.back_bone(inp) loss = self.criterion(out, label) out = torch.round(out).to(int).to('cpu') label = label.to(int).to('cpu') self.train_metrics(out, label) self.log("train_loss", loss.item()) return loss def validation_step(self, val_batch, batch_idx): inp, label = val_batch out = self.back_bone(inp) loss = self.criterion(out, label) out = torch.round(out).to(int).to('cpu') label = label.to(int).to('cpu') self.val_metrics(out, label) self.log("val_loss", loss.item(), prog_bar=True) return loss def on_train_epoch_end(self, outputs): metrics = self.train_metrics.compute() self.logger.experiment.add_scalars('Train', metrics) def on_validation_epoch_end(self): metrics = self.val_metrics.compute() self.logger.experiment.add_scalars('Validation', metrics)
class ImageClassification(BaseModel): """ Model for image classification. This is a configurable class composed by a backbone (see solarnet.models.backbone.py) and a classifier. It is also a LightningModule and nn.Module. """ def __init__( self, n_channel: int = 1, n_class: int = 2, learning_rate: float = 1e-4, class_weight: List[float] = None, backbone: Union[str, nn.Module] = "simple-cnn", backbone_output_size: int = 0, n_hidden: int = 512, dropout: float = 0.2, lr_scheduler: bool = False, lr_scheduler_warmup_steps: int = 100, lr_scheduler_total_steps: int = 0, **kwargs, ): super().__init__() self.save_hyperparameters() if isinstance(backbone, str): self.backbone, backbone_output_size = get_backbone( backbone, channels=n_channel, dropout=dropout, **kwargs, ) self.classifier = Classifier(backbone_output_size, n_class, n_hidden, dropout) if class_weight is not None: class_weight = torch.tensor(class_weight, dtype=torch.float) self.loss_fn = nn.CrossEntropyLoss(weight=class_weight) self.train_accuracy = Accuracy() self.val_accuracy = Accuracy() self.test_metrics = MetricCollection([ Accuracy(), F1(num_classes=self.hparams.n_class, average="macro"), Recall(num_classes=self.hparams.n_class, average="macro"), # balanced acc. StatScores( num_classes=self.hparams.n_class if self.hparams.n_class > 2 else 1, reduce="micro", multiclass=self.hparams.n_class > 2, ), ]) @property def backbone_name(self) -> str: if isinstance(self.hparams.backbone, str): return self.hparams.backbone else: return type(self.hparams.backbone).__name__ @property def output_size(self) -> int: return self.hparams.n_class @auto_move_data def forward(self, image): return self.classifier(self.backbone(image)) def predict_step(self, batch, batch_idx: int, dataloader_idx: int = None): image, _ = batch return self(image) def training_step(self, batch, batch_id): return self.step(batch, step_type="train") def validation_step(self, batch, batch_id): return self.step(batch, step_type="val") def step(self, batch, step_type: str): image, y = batch y_pred = self(image) loss = self.loss_fn(y_pred, y) self.log(f"{step_type}_loss", loss, prog_bar=True, sync_dist=True) # Compute accuracy y_pred = F.softmax(y_pred, dim=1) self.__getattr__(f"{step_type}_accuracy")(y_pred, y) self.log(f"{step_type}_accuracy", self.__getattr__(f"{step_type}_accuracy"), on_step=False, on_epoch=True) return loss def test_step(self, batch, batch_idx): image, y = batch y_pred = self(image) y_pred = F.softmax(y_pred, dim=1) self.test_metrics(y_pred, y) def test_epoch_end(self, outs): test_metrics = self.test_metrics.compute() tp, fp, tn, fn, _ = test_metrics.pop("StatScores") self.log("test_tp", tp) self.log("test_fp", fp) self.log("test_tn", tn) self.log("test_fn", fn) for key, value in test_metrics.items(): self.log(f"test_{key.lower()}", value) def configure_optimizers(self): logger.info(f"configure_optimizers lr={self.hparams.learning_rate}") optimizer = optim.Adam( filter(lambda p: p.requires_grad, self.parameters()), lr=self.hparams.learning_rate, ) if not self.hparams.lr_scheduler: return optimizer scheduler = optim.lr_scheduler.LambdaLR( optimizer, linear_warmup_decay(self.hparams.lr_scheduler_warmup_steps, self.hparams.lr_scheduler_total_steps, cosine=True), ) return ({ "optimizer": optimizer, "lr_scheduler": { "scheduler": scheduler, "interval": "step", "frequency": 1, }, }, )
class TrainingModule(pl.LightningModule): def __init__(self, tagger: LstmTagger): super().__init__() self.tagger = tagger self.train_metrics = MetricCollection([Precision(), Recall(), TopkAccuracy(1)]) self.val_metrics = MetricCollection([Precision(), Recall(), TopkAccuracy(1)]) self.softmax = torch.nn.Softmax(dim=-1) self.celoss = SequenceCrossEntropyLoss(reduction="batch-mean", pad_idx=2) def training_step(self, batch, batch_idx): reports, target, masks = batch mask = torch.cat(masks, dim=1) if self.tagger.with_crf: emissions = torch.cat([self.tagger.calc_emissions(report, mask) for report, mask in zip(reports, masks)], dim=1) loss = -self.tagger.crf(emissions, target, mask) else: scores = self.tagger.forward(reports, masks) loss = self.celoss(scores, target) with torch.no_grad(): scores = self.tagger.forward(reports, masks) preds = scores.argmax(dim=-1) scores = self.softmax(scores) self.train_metrics.update(preds, target, mask, scores=scores) self.log("train_loss", loss) return loss def validation_step(self, batch, *args): reports, target, masks = batch mask = torch.cat(masks, dim=1) if self.tagger.with_crf: emissions = torch.cat([self.tagger.calc_emissions(report, mask) for report, mask in zip(reports, masks)], dim=1) loss = -self.tagger.crf(emissions, target, mask) else: scores = self.tagger.forward(reports, masks) loss = self.celoss(scores, target) with torch.no_grad(): scores = self.tagger.forward(reports, masks) preds = scores.argmax(dim=-1) scores = self.softmax(scores) self.val_metrics.update(preds, target, mask, scores=scores) return loss def validation_epoch_end(self, outputs: List[Any]) -> None: super().validation_epoch_end(outputs) self.log("val_metrics", self.val_metrics.compute()) print(self.val_metrics.compute()) self.val_metrics.reset() def training_epoch_end(self, outputs: List[Any]) -> None: super().training_epoch_end(outputs) self.log("train_metrics", self.train_metrics.compute()) self.train_metrics.reset() def configure_optimizers(self): return Adam(self.parameters(), lr=1e-4, weight_decay=1e-5)
class DeepAnalyze(pl.LightningModule): def __init__(self, feature_size, lstm_hidden_size, lstm_num_layers, n_tags, max_len): super().__init__() self.padding = 0 self.bi_listm = nn.LSTM(feature_size, lstm_hidden_size, num_layers=lstm_num_layers, bidirectional=True) self.attention = DeepAnalyzeAttention(lstm_hidden_size * 2, n_tags, max_len) self.crf = CRF(n_tags) self.lstm_dropout = nn.Dropout(0.25) self.train_metrics = MetricCollection( [Precision(), Recall(), TopkAccuracy(3)]) self.val_metrics = MetricCollection( [Precision(), Recall(), TopkAccuracy(3)]) def forward(self, inputs, mask): seq_len, batch_size = mask.shape x, _ = self.bi_listm(inputs) x = self.lstm_dropout(x) x = self.attention(x, mask) preds = self.crf.decode(x, mask) preds = [pred + [0] * (seq_len - len(pred)) for pred in preds] preds = torch.tensor(preds).transpose(0, 1).to(inputs.device) return preds def training_step(self, batch, batch_idx): inputs, labels, mask = batch x, _ = self.bi_listm(inputs) x = self.lstm_dropout(x) emissions = self.attention(x, mask) loss = -self.crf(emissions, labels, mask) with torch.no_grad(): preds = self.forward(inputs, mask) self.train_metrics.update(preds, labels, mask, scores=get_label_scores( self.crf, emissions, preds, mask)) self.log("train_loss", loss) return loss def validation_step(self, batch, *args): inputs, labels, mask = batch x, _ = self.bi_listm(inputs) x = self.lstm_dropout(x) emissions = self.attention(x, mask) loss = -self.crf(emissions, labels, mask) with torch.no_grad(): preds = self.forward(inputs, mask) self.val_metrics.update(preds, labels, mask, scores=get_label_scores( self.crf, emissions, preds, mask)) return loss def validation_epoch_end(self, outputs: List[Any]) -> None: super().validation_epoch_end(outputs) self.log("val_metrics", self.val_metrics.compute()) self.val_metrics.reset() def training_epoch_end(self, outputs: List[Any]) -> None: super().training_epoch_end(outputs) self.log("train_metrics", self.train_metrics.compute()) self.train_metrics.reset() def configure_optimizers(self): return Adam(self.parameters(), lr=1e-3)