-
Notifications
You must be signed in to change notification settings - Fork 4
/
datasets.py
executable file
·89 lines (65 loc) · 2.82 KB
/
datasets.py
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
84
85
86
87
88
89
import numpy as np
import torch
from torch.utils.data import random_split, DataLoader
from torchvision.transforms import Lambda, ToTensor
from datasets.flatland_ds import ForwardVAEDS
from datasets.dsprites import PairSprites
def sprites_transforms(_):
return ToTensor(), ToTensor()
def forward_ds_transforms(_):
lam = lambda x: torch.from_numpy(np.array(x)).float()
return Lambda(lam), Lambda(lam)
transforms = {
'flatland': forward_ds_transforms,
'dsprites': sprites_transforms,
}
def split(func): # Splits a dataset into a train and val set
def splitter(args):
ds = func(args)
lengths = int(len(ds) * (1 - args.split)), int(len(ds)) - int(len(ds) * (1 - args.split))
train_ds, val_ds = random_split(ds, lengths) if args.split > 0 else (ds, ds)
return train_ds, val_ds
return splitter
def fix_data_path(func): # Sets the datapath to that in _default_paths if it is None
def fixer(args):
args.data_path = args.data_path if args.data_path is not None else _default_paths[args.dataset]
return func(args)
return fixer
def set_to_loader(trainds, valds, args):
trainloader = DataLoader(trainds, batch_size=args.batch_size, num_workers=7, shuffle=args.shuffle, drop_last=False,
pin_memory=True)
valloader = DataLoader(valds, batch_size=args.batch_size, num_workers=7, shuffle=False, drop_last=False,
pin_memory=True)
return trainloader, valloader
@split
@fix_data_path
def sprites(args):
train_transform, test_transform = transforms[args.dataset](args)
output_targets = True if args.model in ['forward', 'rgrvae'] else False
ds = PairSprites(args.data_path, download=False, transform=train_transform, wrapping=True, offset=args.offset,
noise_name=args.noise_name, output_targets=output_targets)
return ds
@split
@fix_data_path
def forward_vae_ds(args):
import os
train_transform, test_transform = transforms[args.dataset](args)
output_targets = True if args.model in ['forward', 'rgrvae'] else False
mean_channels = True
images_path = os.path.join(args.data_path, 'inputs.npy')
actions_path = os.path.join(args.data_path, 'actions.npy')
ds = ForwardVAEDS(images_path, actions_path, transforms=train_transform, output_targets=output_targets,
mean_channels=mean_channels, num_steps=args.offset, noise_name=args.noise_name)
return ds
_default_paths = {
'flatland': '/home/matt/PycharmProjects/PyTorchCode/NeurIPS19-SBDRL/code/learn_4_dim_linear_disentangled_representation/flatland/flat_game',
'dsprites': '',
}
datasets = {
'flatland': forward_vae_ds,
'dsprites': sprites,
}
dataset_meta = {
'flatland': {'nc': 1, 'factors': 2, 'max_classes': 40},
'dsprites': {'nc': 1, 'factors': 5, 'max_classes': 40},
}