autodl-projects/xautodl/datasets/synthetic_env.py

124 lines
3.7 KiB
Python
Raw Normal View History

2021-04-22 17:08:43 +02:00
import math
2021-05-13 15:33:34 +02:00
import random
2021-04-26 14:16:38 +02:00
from typing import List, Optional, Dict
2021-04-22 17:08:43 +02:00
import torch
import torch.utils.data as data
2021-05-13 15:33:34 +02:00
def is_list_tuple(x):
return isinstance(x, (tuple, list))
def zip_sequence(sequence):
def _combine(*alist):
if is_list_tuple(alist[0]):
return [_combine(*xlist) for xlist in zip(*alist)]
else:
return torch.cat(alist, dim=0)
def unsqueeze(a):
if is_list_tuple(a):
return [unsqueeze(x) for x in a]
else:
return a.unsqueeze(dim=0)
with torch.no_grad():
sequence = [unsqueeze(a) for a in sequence]
return _combine(*sequence)
2021-04-26 14:16:38 +02:00
class SyntheticDEnv(data.Dataset):
2021-04-22 17:08:43 +02:00
"""The synethtic dynamic environment."""
def __init__(
self,
2021-05-24 07:06:10 +02:00
data_generator,
oracle_map,
time_generator,
2021-04-22 17:08:43 +02:00
num_per_task: int = 5000,
2021-05-24 07:06:10 +02:00
noise: float = 0.1,
2021-04-22 17:08:43 +02:00
):
2021-05-24 07:06:10 +02:00
self._data_generator = data_generator
self._time_generator = time_generator
self._oracle_map = oracle_map
2021-04-22 17:08:43 +02:00
self._num_per_task = num_per_task
2021-05-24 07:06:10 +02:00
self._noise = noise
2021-05-22 11:36:09 +02:00
2021-04-28 17:56:25 +02:00
@property
def min_timestamp(self):
2021-05-24 07:06:10 +02:00
return self._time_generator.min_timestamp
2021-04-28 17:56:25 +02:00
@property
def max_timestamp(self):
2021-05-24 07:06:10 +02:00
return self._time_generator.max_timestamp
2021-04-28 17:56:25 +02:00
2021-05-10 08:14:06 +02:00
@property
2021-05-24 07:06:10 +02:00
def time_interval(self):
return self._time_generator.interval
@property
def mode(self):
return self._time_generator.mode
2021-05-10 08:14:06 +02:00
2021-05-23 21:14:12 +02:00
def random_timestamp(self, min_timestamp=None, max_timestamp=None):
if min_timestamp is None:
min_timestamp = self.min_timestamp
if max_timestamp is None:
max_timestamp = self.max_timestamp
return random.random() * (max_timestamp - min_timestamp) + min_timestamp
2021-05-13 15:33:34 +02:00
2021-05-13 09:32:44 +02:00
def get_timestamp(self, index):
2021-05-13 15:33:34 +02:00
if index is None:
timestamps = []
2021-05-24 07:06:10 +02:00
for index in range(len(self._time_generator)):
timestamps.append(self._time_generator[index][1])
2021-05-13 15:33:34 +02:00
return tuple(timestamps)
else:
2021-05-24 07:06:10 +02:00
index, timestamp = self._time_generator[index]
2021-05-13 15:33:34 +02:00
return timestamp
2021-05-13 09:32:44 +02:00
2021-04-22 17:08:43 +02:00
def __iter__(self):
self._iter_num = 0
return self
def __next__(self):
if self._iter_num >= len(self):
raise StopIteration
self._iter_num += 1
return self.__getitem__(self._iter_num - 1)
def __getitem__(self, index):
assert 0 <= index < len(self), "{:} is not in [0, {:})".format(index, len(self))
2021-05-24 07:06:10 +02:00
index, timestamp = self._time_generator[index]
2021-05-23 21:14:12 +02:00
return self.__call__(timestamp)
2021-05-22 11:36:09 +02:00
def seq_call(self, timestamps):
with torch.no_grad():
if isinstance(timestamps, torch.Tensor):
timestamps = timestamps.cpu().tolist()
2021-05-13 15:33:34 +02:00
xdata = [self.__call__(timestamp) for timestamp in timestamps]
return zip_sequence(xdata)
2021-05-10 08:14:06 +02:00
def __call__(self, timestamp):
2021-05-24 07:06:10 +02:00
dataset = self._data_generator(timestamp, self._num_per_task)
targets = self._oracle_map.noise_call(dataset, timestamp, self._noise)
return torch.Tensor([timestamp]), (
torch.Tensor(dataset),
torch.Tensor(targets),
2021-04-22 17:08:43 +02:00
)
def __len__(self):
2021-05-24 07:06:10 +02:00
return len(self._time_generator)
2021-04-22 17:08:43 +02:00
def __repr__(self):
2021-05-17 09:39:24 +02:00
return "{name}({cur_num:}/{total} elements, ndim={ndim}, num_per_task={num_per_task}, range=[{xrange_min:.5f}~{xrange_max:.5f}], mode={mode})".format(
2021-04-22 17:08:43 +02:00
name=self.__class__.__name__,
cur_num=len(self),
2021-05-24 07:06:10 +02:00
total=len(self._time_generator),
2021-05-24 07:14:39 +02:00
ndim=self._data_generator.ndim,
2021-04-22 17:08:43 +02:00
num_per_task=self._num_per_task,
2021-05-17 09:39:24 +02:00
xrange_min=self.min_timestamp,
xrange_max=self.max_timestamp,
2021-05-24 07:06:10 +02:00
mode=self.mode,
2021-04-22 17:08:43 +02:00
)