def __init__(self, num_hidden, num_input, num_batch=1, w_init=XavierNormal(), b_init=ZeroInit(), has_bias=True, s_init=Uniform(), train_mask=None, **kwargs): super(LinearReadout, self).__init__(size=num_hidden, **kwargs) # parameters self.w_init = w_init self.b_init = b_init self.s_init = s_init self.num_input = num_input self.has_bias = has_bias # weights self.w = math.TrainVar(w_init((num_input, self.num))) if has_bias: self.b = math.TrainVar(b_init((self.num, ))) if train_mask is not None: assert train_mask.shape == self.w.shape self.train_mask = train_mask # variables self.s = math.Variable(self.s_init((num_batch, self.num)))
def __init__(self, num_hidden, num_input, num_batch, wx=Orthogonal(), wh=Orthogonal(), b=ZeroInit(), h=ZeroInit(), **kwargs): super(GRU, self).__init__(num_hidden, num_input, **kwargs) self.has_bias = True # variables self.h = bm.Variable(self.get_param(h, (num_batch, self.num_hidden))) # weights wxs = self.get_param(wx, (num_input * 3, num_hidden)) self.w_iz = bm.TrainVar(wxs[:num_input]) self.w_ir = bm.TrainVar(wxs[num_input:num_input * 2]) self.w_ia = bm.TrainVar(wxs[num_input * 2:]) whs = self.get_param(wh, (num_hidden * 3, num_hidden)) self.w_hz = bm.TrainVar(whs[:num_hidden]) self.w_hr = bm.TrainVar(whs[num_hidden:num_hidden * 2]) self.w_ha = bm.TrainVar(whs[num_hidden * 2:]) bs = self.get_param(b, (num_hidden * 3, )) self.bz = bm.TrainVar(bs[:num_hidden]) self.br = bm.TrainVar(bs[num_hidden:num_hidden * 2]) self.ba = bm.TrainVar(bs[num_hidden * 2:])
def get_param(param, size): if param is None: return None if callable(param): return bm.TrainVar(param(size)) if isinstance(param, onp.ndarray): assert param.shape == size return bm.TrainVar(bm.asarray(param)) if isinstance(param, (bm.JaxArray, jnp.ndarray)): return bm.TrainVar(param) raise ValueError
def __init__(self, num_input, num_hidden, num_output, num_batch, dt=None, e_ratio=0.8, sigma_rec=0., seed=None, w_ir=bp.init.KaimingUniform(scale=1.), w_rr=bp.init.KaimingUniform(scale=1.), w_ro=bp.init.KaimingUniform(scale=1.)): super(RNN, self).__init__() # parameters self.tau = 100 self.num_batch = num_batch self.num_input = num_input self.num_hidden = num_hidden self.num_output = num_output self.e_size = int(num_hidden * e_ratio) self.i_size = num_hidden - self.e_size if dt is None: self.alpha = 1 else: self.alpha = dt / self.tau self.sigma_rec = (2 * self.alpha)**0.5 * sigma_rec # Recurrent noise self.rng = bm.random.RandomState(seed=seed) # hidden mask mask = np.tile([1] * self.e_size + [-1] * self.i_size, (num_hidden, 1)) np.fill_diagonal(mask, 0) self.mask = bm.asarray(mask, dtype=bm.float_) # input weight self.w_ir = self.get_param(w_ir, (num_input, num_hidden)) # recurrent weight bound = 1 / num_hidden**0.5 self.w_rr = self.get_param(w_rr, (num_hidden, num_hidden)) self.w_rr[:, :self.e_size] /= (self.e_size / self.i_size) self.b_rr = bm.TrainVar(self.rng.uniform(-bound, bound, num_hidden)) # readout weight bound = 1 / self.e_size**0.5 self.w_ro = self.get_param(w_ro, (self.e_size, num_output)) self.b_ro = bm.TrainVar(self.rng.uniform(-bound, bound, num_output)) # variables self.h = bm.Variable(bm.zeros((num_batch, num_hidden))) self.o = bm.Variable(bm.zeros((num_batch, num_output)))
def __init__(self, num_hidden, num_input, num_batch, h=Uniform(), w=XavierNormal(), b=ZeroInit(), **kwargs): super(VanillaRNN, self).__init__(num_hidden, num_input, **kwargs) # variables self.h = bm.Variable(self.get_param(h, (num_batch, self.num_hidden))) # weights ws = self.get_param(w, (num_input + num_hidden, num_hidden)) self.w_ir = bm.TrainVar(ws[:num_input]) self.w_rr = bm.TrainVar(ws[num_input:]) self.b = self.get_param(b, (num_hidden, ))
def __init__(self, num_input, num_hidden, num_output, num_batch, dt=None, seed=None, w_ir=bp.init.KaimingNormal(scale=1.), w_rr=bp.init.KaimingNormal(scale=1.), w_ro=bp.init.KaimingNormal(scale=1.)): super(RNN, self).__init__() # parameters self.tau = 100 self.num_batch = num_batch self.num_input = num_input self.num_hidden = num_hidden self.num_output = num_output if dt is None: self.alpha = 1 else: self.alpha = dt / self.tau self.rng = bm.random.RandomState(seed=seed) # input weight self.w_ir = self.get_param(w_ir, (num_input, num_hidden)) # recurrent weight bound = 1 / num_hidden**0.5 self.w_rr = self.get_param(w_rr, (num_hidden, num_hidden)) self.b_rr = bm.TrainVar(self.rng.uniform(-bound, bound, num_hidden)) # readout weight self.w_ro = self.get_param(w_ro, (num_hidden, num_output)) self.b_ro = bm.TrainVar(self.rng.uniform(-bound, bound, num_output)) # variables self.h = bm.Variable(bm.zeros((num_batch, num_hidden))) self.o = bm.Variable(bm.zeros((num_batch, num_output)))
def __init__(self): super(Test, self).__init__() self.a = bm.TrainVar(bm.ones(10)) self.b = bm.TrainVar(bm.random.randn(10)) self.c = bm.TrainVar(bm.random.uniform(size=10))