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)
