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)
