Spaces:
Sleeping
Sleeping
File size: 4,792 Bytes
165ee00 |
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 90 91 92 93 94 95 96 97 98 99 |
from abc import ABCMeta, abstractmethod
from typing import Set, Optional
from dataclasses import dataclass, fields
import torch
from torch.utils.data import DataLoader
@dataclass
class Batch:
"""
A batch of data, with non-optional x, y, and target_y attributes.
All other attributes are optional.
If you want to add an attribute for testing only, you can just assign it after creation like:
```
batch = Batch(x=x, y=y, target_y=target_y)
batch.test_attribute = test_attribute
```
"""
# Required entries
x: torch.Tensor
y: torch.Tensor
target_y: torch.Tensor
# Optional Batch Entries
style: Optional[torch.Tensor] = None
style_hyperparameter_values: Optional[torch.Tensor] = None
single_eval_pos: Optional[torch.Tensor] = None
causal_model_dag: Optional[object] = None
mean_prediction: Optional[bool] = None # this controls whether to do mean prediction in bar_distribution for nonmyopic BO
def other_filled_attributes(self, set_of_attributes: Set[str] = frozenset(('x', 'y', 'target_y'))):
return [f.name for f in fields(self)
if f.name not in set_of_attributes and
getattr(self, f.name) is not None]
def safe_merge_batches_in_batch_dim(*batches, ignore_attributes=[]):
"""
Merge all supported non-None fields in a pre-specified (general) way,
e.g. mutliple batch.x are concatenated in the batch dimension.
:param ignore_attributes: attributes to remove from the merged batch, treated as if they were None.
:return:
"""
not_none_fields = [f.name for f in fields(batches[0]) if f.name not in ignore_attributes and getattr(batches[0], f.name) is not None]
assert all([set(not_none_fields) == set([f.name for f in fields(b) if f.name not in ignore_attributes and getattr(b, f.name) is not None]) for b in batches]), 'All batches must have the same fields!'
merge_funcs = {
'x': lambda xs: torch.cat(xs, 1),
'y': lambda ys: torch.cat(ys, 1),
'target_y': lambda target_ys: torch.cat(target_ys, 1),
'style': lambda styles: torch.cat(styles, 0),
}
assert all(f in merge_funcs for f in not_none_fields), f'Unknown fields encountered in `safe_merge_batches_in_batch_dim`.'
return Batch(**{f: merge_funcs[f]([getattr(batch, f) for batch in batches]) for f in not_none_fields})
def merge_batches(*batches, ignore_attributes=[]):
assert False, "TODO: isn't this broken!? because catting in dim 0 seems wrong!?"
def merge_attribute(attr_name, batch_sizes):
attr = [getattr(batch, attr_name) for batch in batches]
if type(attr[0]) is list:
def make_list(sublist, i):
if sublist is None:
return [None for _ in range(batch_sizes[i])]
return sublist
return sum([make_list(sublist, i) for i, sublist in enumerate(attr)], [])
elif type(attr[0]) is torch.Tensor:
return torch.cat(attr, 0)
else:
assert all(a is None for a in attr), f'Unknown type encountered in `merge_batches`.'\
f'To ignore this, please add `{attr}` to the `ignore_attributes`.'\
f'The following values are the problem: {attr_name}.'
return None
batch_sizes = [batch.x.shape[0] for batch in batches]
return Batch(**{f.name: merge_attribute(f.name, batch_sizes) for f in fields(batches[0]) if f.name not in ignore_attributes})
class PriorDataLoader(DataLoader, metaclass=ABCMeta):
@abstractmethod
def __init__(self, num_steps, batch_size, eval_pos_seq_len_sampler, seq_len_maximum, device, **kwargs):
"""
:param num_steps: int, first argument, the number of steps to take per epoch, i.e. iteration of the DataLoader
:param batch_size: int, number of datasets per batch
:param eval_pos_seq_len_sampler: callable, it takes no arguments and returns a tuple (single eval pos, bptt)
:param kwargs: for future compatibility it is good to have a final all catch, as new kwargs might be introduced
"""
pass
# A class or object variable `num_features`: int
# Optional: `validate` function that accepts a transformer model
# The DataLoader iter should return batches of the form ([style], x, y), target_y, single_eval_pos
# We follow sequence len (s) first, batch size (b) second. So x: (s,b,num_features), y,target_y: (s,b)
# and style: Optional[(b,num_style_params)], style can be omitted or set to None, if it is not intended to be used.
# For more references, see `priors/utils.py` for a pretty general implementation of a DataLoader
# and `train.py` for the only call of it.
|