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 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)
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 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)