| 70 | |
| 71 | # define the LightningModule |
| 72 | class LocationEncoder(pl.LightningModule): |
| 73 | def __init__(self, positional_encoding_name, neural_network_name, hparams): |
| 74 | super().__init__() |
| 75 | |
| 76 | self.learning_rate = hparams["optimizer"]["lr"] |
| 77 | self.weight_decay = hparams["optimizer"]["wd"] |
| 78 | self.regression = get_param(hparams, "regression") |
| 79 | |
| 80 | self.loss_fn = get_loss_fn(presence_only=get_param(hparams, "presence_only_loss"), |
| 81 | loss_weight=get_param(hparams, "loss_weight"), |
| 82 | regression=self.regression) |
| 83 | |
| 84 | self.positional_encoder = get_positional_encoding( |
| 85 | positional_encoding_name, hparams |
| 86 | ) |
| 87 | self.neural_network = get_neural_network( |
| 88 | neural_network_name, |
| 89 | input_dim=self.positional_encoder.embedding_dim, |
| 90 | hparams=hparams |
| 91 | ) |
| 92 | |
| 93 | # this enables LocationEncoder.load_from_checkpoint(path) |
| 94 | self.save_hyperparameters() |
| 95 | |
| 96 | def common_step(self, batch, batch_idx): |
| 97 | lonlats, label = batch |
| 98 | return self.loss_fn(self, lonlats, label) |
| 99 | |
| 100 | def forward(self, lonlats): |
| 101 | embedding = self.positional_encoder(lonlats) |
| 102 | return self.neural_network(embedding) |
| 103 | |
| 104 | def training_step(self, batch, batch_idx): |
| 105 | loss = self.common_step(batch, batch_idx) |
| 106 | return loss |
| 107 | |
| 108 | def validation_step(self, batch, batch_idx): |
| 109 | loss = self.common_step(batch, batch_idx) |
| 110 | self.log("val_loss", loss, on_step=False, on_epoch=True) |
| 111 | return {"val_loss":loss} |
| 112 | |
| 113 | def predict_step(self, batch, batch_idx): |
| 114 | lonlats, label = batch |
| 115 | prediction_logits = self.forward(lonlats) |
| 116 | return prediction_logits, lonlats, label |
| 117 | |
| 118 | def test_step(self, batch, batch_idx): |
| 119 | lonlats, label = batch |
| 120 | prediction_logits = self.forward(lonlats) |
| 121 | |
| 122 | loss = self.loss_fn(self, lonlats, label) |
| 123 | |
| 124 | # check if binary |
| 125 | non_binary_task = self.regression |
| 126 | if (prediction_logits.size(1) == 1) and not (non_binary_task): |
| 127 | y_pred = (prediction_logits.squeeze() > 0).cpu() |
| 128 | average = "binary" |
| 129 | elif self.regression: |