BioLCNet/dataset.py at main · Singular-Brain/BioLCNet · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

import torch
from bindsnet.encoding import PoissonEncoder

from bindsnet.datasets import MNIST

import os

from torchvision import transforms


class ClassSelector(torch.utils.data.sampler.Sampler):
    """Select target classes from the dataset"""
    def __init__(self, target_classes, data_source, mask = None):
        if mask is not None:
            self.mask = mask
        else:
            self.mask = torch.tensor([1 if data_source[i]['label'] in target_classes else 0 for i in range(len(data_source))])
        self.data_source = data_source

    def __iter__(self):
        return iter([i.item() for i in torch.nonzero(self.mask)])

    def __len__(self):
        return len(self.data_source)


# Load MNIST data.
def load_datasets(data_hparams, target_classes=None, mask=None, mask_test=None):
    dataset = MNIST(
        PoissonEncoder(time=data_hparams['time'], dt=1),
        None,
        root=os.path.join("MNIST"),
        download=True,
        transform=transforms.Compose(
            [transforms.ToTensor(),
            transforms.Lambda(lambda x: (x * data_hparams['intensity'])),
            transforms.CenterCrop(data_hparams['crop_size'])]
        ),
    )

    # Create a dataloader to iterate and batch data
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1,
                                            sampler = ClassSelector(
                                                    target_classes = target_classes,
                                                    data_source = dataset,
                                                    mask = mask,
                                                    ) if target_classes else None
                                            )

    val_loader = torch.utils.data.DataLoader(dataset, batch_size=1,
                                            sampler = ClassSelector(
                                                    target_classes = target_classes,
                                                    data_source = dataset,
                                                    mask = mask,
                                                    ) if target_classes else None
                                            )

    # Load test dataset
    test_dataset = MNIST(
        PoissonEncoder(time=data_hparams['time'], dt=1),
        None,
        root=os.path.join("MNIST"),
        download=True,
        train=False,
        transform=transforms.Compose(
            [transforms.ToTensor(),
            transforms.Lambda(lambda x: (x * data_hparams['intensity'])),
            transforms.CenterCrop(data_hparams['crop_size'])]
        ),
    )


    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1,
                                        sampler = ClassSelector(
                                                target_classes = target_classes,
                                                data_source = test_dataset,
                                                mask = mask_test,
                                                ) if target_classes else None
                                        )

    return dataloader, val_loader, test_loader