From bf0b5ba3dc630d2aa5c113b44fdafbbb28512bba Mon Sep 17 00:00:00 2001 From: "Daniel.Frisinghelli" <daniel.frisinghelli@eurac.edu> Date: Wed, 3 Mar 2021 14:24:22 +0100 Subject: [PATCH] Computing cm on GPU requires rewriting sklearn cm. Aborting. --- pysegcnn/core/trainer.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py index dac20d2..12cdafb 100644 --- a/pysegcnn/core/trainer.py +++ b/pysegcnn/core/trainer.py @@ -2731,18 +2731,12 @@ 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) - labels_gpu = torch.Tensor( - np.asarray(list(self.use_labels.keys()))).to(self.device) - - # compute confusion matrix - conf_mat = confusion_matrix(y_true, y_pred, labels=labels_gpu) + 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.cpu().numpy() + output['cm'] = conf_mat # plot confusion matrix plot_confusion_matrix( -- GitLab