Spaces:
Sleeping
Sleeping
| from abc import ABCMeta, abstractmethod | |
| from typing import Set, Optional | |
| from dataclasses import dataclass, fields | |
| import torch | |
| from torch.utils.data import DataLoader | |
| 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): | |
| 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. | |