def __init__(self, data_name):
super(OnlineToyDataset, self).__init__(2)
self.data_name = data_name
self.rng = np.random.RandomState()
rng = np.random.RandomState(1)
samples = inf_train_gen(self.data_name, rng, 5000)
self.f_scale = np.max(np.abs(samples)) + 1
self.int_scale = 2 ** 15 / (self.f_scale + 1)
print('f_scale,', self.f_scale, 'int_scale,', self.int_scale)
def gen_batch(self, batch_size):
return inf_train_gen(self.data_name, self.rng, batch_size)