From 00ef129b2558bde592898bc3f64970b4e51fe4f3 Mon Sep 17 00:00:00 2001
From: "Daniel.Frisinghelli" <daniel.frisinghelli@eurac.edu>
Date: Wed, 3 Mar 2021 14:13:40 +0100
Subject: [PATCH] Compute confusion matrix on GPU, if available.
---
pysegcnn/core/trainer.py | 8 +++++++-
1 file changed, 7 insertions(+), 1 deletion(-)
diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index 12cdafb..a5d3506 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2731,12 +2731,18 @@ class NetworkInference(BaseConfig):
# calculate confusion matrix
LOGGER.info('Computing confusion matrix ...')
+
+ # move predictions and labels to GPU if available
+ y_true = torch.Tensor(y_true).to(self.device)
+ y_pred = torch.Tensor(y_pred).to(self.device)
+
+ # compute confusion matrix
conf_mat = confusion_matrix(
y_true, y_pred,
labels=np.asarray(list(self.use_labels.keys())))
# add confusion matrix to model output
- output['cm'] = conf_mat
+ output['cm'] = conf_mat.cpu().numpy()
# plot confusion matrix
plot_confusion_matrix(
--
GitLab