| 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: |
| 130 | y_pred = prediction_logits.cpu() |
| 131 | else: # take argmax |
| 132 | y_pred = prediction_logits.argmax(-1).cpu() |
| 133 | average = "macro" |
| 134 | |
| 135 | self.log("test_loss", loss, on_step=False, on_epoch=True) |
| 136 | if self.regression: |
| 137 | MAE = mean_absolute_error(y_true=label.cpu(), y_pred = y_pred) |
| 138 | self.log("test_MAE", MAE, on_step=False, on_epoch=True) |
| 139 | |
| 140 | test_results = {"test_loss":loss, |
| 141 | "test_MAE":MAE} |
| 142 | |
| 143 | else: |
| 144 | accuracy = accuracy_score(y_true=label.cpu(), y_pred= y_pred).astype("float32") |
| 145 | IoU = jaccard_score(y_true=label.cpu(), y_pred = y_pred, average=average).astype("float32") |
| 146 | self.log("test_accuracy", accuracy, on_step=False, on_epoch=True) |
| 147 | self.log("test_IoU", IoU, on_step=False, on_epoch=True) |
| 148 | |
| 149 | test_results = {"test_loss":loss, |
| 150 | "test_accuracy":accuracy} |
| 151 | |
| 152 | return test_results |
| 153 | |
| 154 | def configure_optimizers(self): |
| 155 | optimizer = optim.Adam([{"params": self.neural_network.parameters()}, |