Spaces:
Sleeping
Sleeping
File size: 2,925 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 |
import torch
from .prior import Batch
from ..utils import default_device
@torch.no_grad()
def get_batch(batch_size, seq_len, num_features, get_batch, epoch, device=default_device, hyperparameters={}, **kwargs):
"""
This is not part of the paper, but feel welcome to use this to write a better version of our user prior.
This function assumes that every x is in the range [0, 1].
Style shape is (batch_size, 3*num_features) under the assumption that get_batch returns a batch
with shape (seq_len, batch_size, num_features).
The style is build the following way: [prob_of_feature_1_in_range, range_min_of_feature_1, range_max_of_feature_1, ...]
:param batch_size:
:param seq_len:
:param num_features:
:param get_batch:
:param epoch:
:param device:
:param hyperparameters:
:param kwargs:
:return:
"""
maximize = hyperparameters.get('condition_on_area_maximize', True)
size_range = hyperparameters.get('condition_on_area_size_range', (0.1, 0.5))
distribution = hyperparameters.get('condition_on_area_distribution', 'uniform')
assert distribution in ['uniform']
batch: Batch = get_batch(batch_size=batch_size, seq_len=seq_len,
num_features=num_features, device=device,
hyperparameters=hyperparameters,
epoch=epoch, **kwargs)
assert batch.style is None
d = batch.x.shape[2]
prob_correct = torch.rand(batch_size, d, device=device)
correct_opt = torch.rand(batch_size, d, device=device) < prob_correct
division_size = torch.rand(batch_size, d, device=device) * (size_range[1] - size_range[0]) + size_range[0]
optima = batch.target_y.argmax(0).squeeze() if maximize else batch.target_y.argmin(0).squeeze() # batch_size, d
optima_hints = batch.x[optima, torch.arange(batch_size, device=device)] - division_size/2 + torch.rand(batch_size, d, device=device)*division_size # shape: (batch_size, d)
optima_hints = optima_hints.clamp(0, 1)
optima_division_lower_bound = (optima_hints - division_size/2).clamp(0, 1)
optima_division_upper_bound = (optima_hints + division_size/2).clamp(0, 1)
random_hints = torch.rand(batch_size, d, device=device) - division_size/2 + torch.rand(batch_size, d, device=device)*division_size # shape: (batch_size, d)
random_hints = random_hints.clamp(0, 1)
random_division_lower_bound = (random_hints - division_size/2).clamp(0, 1)
random_division_upper_bound = (random_hints + division_size/2).clamp(0, 1)
lower_bounds = torch.where(correct_opt, optima_division_lower_bound, random_division_lower_bound)
upper_bounds = torch.where(correct_opt, optima_division_upper_bound, random_division_upper_bound)
batch.style = torch.stack([prob_correct, lower_bounds, upper_bounds], 2).view(batch_size, -1) # shape: (batch_size, 3*d)
return batch
|