Bugfix: new torch versions do not support sparse tensors in default_c…

…ollate (#73)

Signed-off-by: Matthias Kümmerer <[email protected]>

pysaliency/saliency_map_conversion_torch.py

@@ -7,7 +7,7 @@
 from .optpy import minimize
 from .saliency_map_models import SaliencyMapModel
 from .torch_utils import GaussianFilterNd, Nonlinearity, zero_grad, log_likelihood
-from .torch_datasets import ImageDataset, ImageDatasetSampler, FixationMaskTransform
+from .torch_datasets import ImageDataset, ImageDatasetSampler, FixationMaskTransform, collate_fn
 
 
 class CenterBias(nn.Module):
@@ -275,6 +275,7 @@ def _optimize_saliency_map_conversion_over_multiple_models_and_datasets(
         batch_sampler=ImageDatasetSampler(dataset, batch_size=batch_size, shuffle=False),
         pin_memory=False,
         num_workers=0,  # doesn't work for sparse tensors yet. Might work soon.
+        collate_fn=collate_fn,
     )
 
     saliency_map_processing = SaliencyMapProcessing(

pysaliency/torch_datasets.py

@@ -3,6 +3,7 @@
 import numpy as np
 import torch
 from boltons.iterutils import chunked
+from torch.utils.data.dataloader import default_collate
 from tqdm import tqdm
 
 from .models import Model
@@ -167,3 +168,14 @@ def __iter__(self):
 
     def __len__(self):
         return int(self.ratio_used * len(self.batches))
+
+
+# we need to extend the defaut collate fn to handle sparse coo tensors
+def collate_fn(batch):
+    result = {}
+    for key in batch[0]:
+        if isinstance(batch[0][key], torch.sparse.Tensor):
+            result[key] = torch.stack([item[key] for item in batch], 0)
+        else:
+            result[key] = default_collate([item[key] for item in batch])
+    return result

tests/test_torch_datasets.py

@@ -11,7 +11,7 @@
     SaliencyMapModelFromDirectory,
     UniformModel
 )
-from pysaliency.torch_datasets import ImageDataset, ImageDatasetSampler, FixationMaskTransform
+from pysaliency.torch_datasets import ImageDataset, ImageDatasetSampler, FixationMaskTransform, collate_fn
 import torch
 
 
@@ -71,6 +71,7 @@ def test_dataset(stimuli, fixations, png_saliency_map_model):
         batch_sampler=ImageDatasetSampler(dataset, batch_size=4, shuffle=False),
         pin_memory=False,
         num_workers=0,  # doesn't work for sparse tensors yet. Might work soon.
+        collate_fn=collate_fn,
     )
 
     count = 0