def __init__(self): self.ps1 = jt.nn.ParameterList( [jt.array([1, 2, 3]), jt.rand(10)]) self.ps2 = jt.nn.ParameterDict({ "aaa": jt.array([1, 2, 3]), "bbb": jt.rand(10) })
def test_fuse_transpose5(self): with jt.profile_scope() as rep: a = jt.rand((10, 11, 6, 7)) c = jt.rand((10, 11, 6, 7)) b = (a + c).fuse_transpose((1, 0, 2, 3)) np.testing.assert_allclose((a.data + c.data).transpose( (1, 0, 2, 3)), b.data) assert len(rep) == 3
def test_fuse_transpose3(self): with jt.profile_scope() as rep: a = jt.rand((10, 11, 12)) c = jt.rand((11, 12, 10)) b = a.fuse_transpose((1, 2, 0)) + c np.testing.assert_allclose( a.data.transpose((1, 2, 0)) + c.data, b.data) assert len(rep) == 3
def test_normal(self): from jittor import init n = 10000 r = 0.155 a = init.gauss([n], "float32", 1, 3) data = a.data assert (np.abs((data < (1 - 3)).mean() - r) < 0.1) assert (np.abs((data < (1)).mean() - 0.5) < 0.1) assert (np.abs((data < (1 + 3)).mean() - (1 - r)) < 0.1) np_res = np.random.normal(1, 0.1, (100, 100)) jt_res = jt.normal(1., 0.1, (100, 100)) assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1) assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1) np_res = torch.normal(torch.arange(1., 10000.), 1) jt_res = jt.normal(jt.arange(1, 10000), 1) assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1) assert (np.abs(np_res.std() - jt_res.data.std()) < 1) np_res = np.random.randn(100, 100) jt_res = jt.randn(100, 100) assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1) assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1) np_res = np.random.rand(100, 100) jt_res = jt.rand(100, 100) assert (np.abs(np_res.mean() - jt_res.data.mean()) < 0.1) assert (np.abs(np_res.std() - jt_res.data.std()) < 0.1)
def sample(self, sample_shape=()): shape = sample_shape + self.probs.shape[:-1] + (1, ) rand = jt.rand(shape) one_hot = jt.logical_and(self.cum_probs_l < rand, rand <= self.cum_probs_r) index = one_hot.index(one_hot.ndim - 1) return (one_hot * index).sum(-1)
def test_cutt(self): a = jt.rand((10, 2)) > 0.5 b = a.transpose() assert (a.data.transpose() == b.data).all() a = jt.zeros((1, 1)) b = a.transpose((1, 0)) b.sync()
def test_conv1d(self): conv1d = jt.nn.Conv1d(10, 20, 5) a = jt.rand((3, 10, 15)) b = conv1d(a) b.sync() assert b.shape == [3, 20, 11] b = jt.nn.Conv1d(10, 20, 5, padding=2)(a) assert b.shape == [3, 20, 15] assert sorted(list(conv1d.state_dict().keys())) == [ 'bias', 'weight' ], conv1d.state_dict().keys()
def test_twobilinear_lstm(self): x = jt.rand(5, 4, 10) rnn1 = nn.LSTM(10, 20, bidirectional=True) out1, _ = rnn1(x) rnn2 = nn.LSTM(40, 20, bidirectional=True) out2, _ = rnn2(out1) target = jt.zeros_like(out2) loss = nn.mse_loss(out2, target) from jittor import optim optimizer = optim.RMSprop(rnn1.parameters()) optimizer.step(loss)
def test_fuse_transpose2(self): with jt.profile_scope() as rep: a = jt.rand((10, 11, 12)) b = (a + 1).fuse_transpose((1, 2, 0)) np.testing.assert_allclose(a.data.transpose((1, 2, 0)) + 1, b.data) assert len(rep) == 3
def test_cutt(self): a = jt.rand((10, 2)) > 0.5 b = a.transpose() assert (a.data.transpose() == b.data).all()
baseWidth=26, scale=4, **kwargs) if pretrained: model.load(jt.load(model_urls['res2net101_v1b_26w_4s'])) return model def res2net50_26w_4s(pretrained=False, **kwargs): model = Res2Net(Bottle2neck, [3, 4, 6, 3], baseWidth=26, scale=4, **kwargs) if pretrained: model.load(jt.load(model_urls['res2net50_v1b_26w_4s'])) return model def res2net101_v1b_26w_4s(pretrained=False, **kwargs): model = Res2Net(Bottle2neck, [3, 4, 23, 3], baseWidth=26, scale=4, **kwargs) if pretrained: model.load(jt.load((model_urls['res2net101_v1b_26w_4s']))) return model if __name__ == '__main__': images = jt.rand(1, 3, 352, 352) model = res2net50_26w_4s(pretrained=False) model = model print(model(images).shape)
def sample(self, sample_shape): tiny = jt.info(self.probs.dtype).tiny u = jt.clamp(jt.rand(sample_shape), min_v=tiny) return (jt.log(u) / (jt.log(-self.probs + 1))).floor()
def summary_string(model, input_size, batch_size=-1, device='cpu', dtypes=None): assert (device in device_list) if device == 'cuda': jt.flags.use_cuda = 1 else: jt.flags.use_cuda = 0 if dtypes == None: dtypes = [jt.float] * len(input_size) summary_str = '' def register_hook(module): def hook(module, input, output): class_name = str(module.__class__).split(".")[-1].split("'")[0] module_idx = len(summary) m_key = "%s-%i" % (class_name, module_idx + 1) summary[m_key] = OrderedDict() summary[m_key]["input_shape"] = list(input[0].size()) summary[m_key]["input_shape"][0] = batch_size if isinstance(output, (list, tuple)): summary[m_key]["output_shape"] = [[-1] + list(o.size())[1:] for o in output] else: summary[m_key]["output_shape"] = list(output.size()) summary[m_key]["output_shape"][0] = batch_size params = 0 if hasattr(module, "weight") and hasattr(module.weight, "size"): params += np.prod( np.array(list(module.weight.size()), dtype=np.int64)) summary[m_key]["trainable"] = module.weight.requires_grad if hasattr(module, "bias") and hasattr(module.bias, "size"): params += np.prod( np.array(list(module.bias.size()), dtype=np.int64)) summary[m_key]["nb_params"] = params if (not isinstance(module, nn.Sequential) and not isinstance(module, nn.ModuleList)): hooks.append(module.register_forward_hook(hook)) # multiple inputs to the network if isinstance(input_size, tuple): input_size = [input_size] # batch_size of 2 for batchnorm x = [jt.rand(2, *in_size).float() for in_size in input_size] # create properties summary = OrderedDict() hooks = [] # register hook model.apply(register_hook) # make a forward pass # print(x.shape) model(*x) # remove these hooks for h in hooks: if h: h.remove() summary_str += "----------------------------------------------------------------" + "\n" line_new = "{:>20} {:>25} {:>15}".format("Layer (type)", "Output Shape", "Param #") summary_str += line_new + "\n" summary_str += "================================================================" + "\n" total_params = 0 total_output = 0 trainable_params = 0 for layer in summary: # input_shape, output_shape, trainable, nb_params line_new = "{:>20} {:>25} {:>15}".format( layer, str(summary[layer]["output_shape"]), "{0:,}".format(summary[layer]["nb_params"]), ) total_params += summary[layer]["nb_params"] total_output += np.prod(summary[layer]["output_shape"]) if "trainable" in summary[layer]: if summary[layer]["trainable"] == True: trainable_params += summary[layer]["nb_params"] summary_str += line_new + "\n" # assume 4 bytes/number (float on cuda). total_input_size = abs( np.prod(sum(input_size, ())) * batch_size * 4. / (1024**2.)) total_output_size = abs(2. * total_output * 4. / (1024**2.)) # x2 for gradients total_params_size = abs(total_params * 4. / (1024**2.)) total_size = total_params_size + total_output_size + total_input_size summary_str += "================================================================" + "\n" summary_str += "Total params: {0:,}".format(total_params) + "\n" summary_str += "Trainable params: {0:,}".format(trainable_params) + "\n" summary_str += "Non-trainable params: {0:,}".format( total_params - trainable_params) + "\n" summary_str += "----------------------------------------------------------------" + "\n" summary_str += "Input size (MB): %0.2f" % total_input_size + "\n" summary_str += "Forward/backward pass size (MB): %0.2f" % total_output_size + "\n" summary_str += "Params size (MB): %0.2f" % total_params_size + "\n" summary_str += "Estimated Total Size (MB): %0.2f" % total_size + "\n" summary_str += "----------------------------------------------------------------" + "\n" # return summary return summary_str, (total_params, trainable_params)
def sample(self, sample_shape): u = jt.rand(sample_shape) return (jt.safe_log(u) / (jt.safe_log(-self.probs + 1))).floor()
def __getitem__(self, k): return jt.rand(2)
def sample(self, sample_shape=[]): shape = sample_shape + self.probs.shape[:-1] + (1, ) rand = jt.rand(shape) one_hot = jt.logical_and(self.cum_probs_l < rand, rand <= self.cum_probs_r).float() return one_hot
def test_param_list(self): ps = jt.nn.ParameterList([jt.array([1, 2, 3]), jt.rand(10)]) assert len(ps.parameters()) == 2 assert list(ps.state_dict().keys()) == ['0', '1'], ps.state_dict().keys()
def test_flip_grad(self): a = jt.rand(10) b = a[::-1] c = b[::-1] d = c.sum() jt.grad(d, [a])
def test_interpolate(self): a = jt.rand(1,3,64,64) b = jt.nn.interpolate(a, scale_factor=0.5) b.sync() assert b.shape == (1,3,32,32)