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)
