def test_lstm_cell(self):
np_h0 = torch.randn(3, 20).numpy()
np_c0 = torch.randn(3, 20).numpy()
t_rnn = tnn.LSTMCell(10, 20)
input = torch.randn(2, 3, 10)
h0 = torch.from_numpy(np_h0)
c0 = torch.from_numpy(np_c0)
t_output = []
for i in range(input.size()[0]):
h0, c0 = t_rnn(input[i], (h0, c0))
t_output.append(h0)
t_output = torch.stack(t_output, dim=0)
j_rnn = nn.LSTMCell(10, 20)
j_rnn.load_state_dict(t_rnn.state_dict())
input = jt.float32(input.numpy())
h0 = jt.float32(np_h0)
c0 = jt.float32(np_c0)
j_output = []
for i in range(input.size()[0]):
h0, c0 = j_rnn(input[i], (h0, c0))
j_output.append(h0)
j_output = jt.stack(j_output, dim=0)
t_output = t_output.detach().numpy()
j_output = j_output.data
assert np.allclose(t_output, j_output, rtol=1e-03, atol=1e-06)
def check_equal_2(t_rnn, j_rnn, input, h0, c0, dev=None):
j_rnn.load_state_dict(t_rnn.state_dict())
if dev:
t_output, (th, tc) = t_rnn(
torch.from_numpy(input).to(dev),
(torch.from_numpy(h0).to(dev), torch.from_numpy(c0).to(dev)))
else:
t_output, (th, tc) = t_rnn(
torch.from_numpy(input).to(dev),
(torch.from_numpy(h0), torch.from_numpy(c0)))
j_output, (jh, jc) = j_rnn(jt.float32(input),
(jt.float32(h0), jt.float32(c0)))
np.testing.assert_allclose(t_output.detach().cpu().numpy(),
j_output.data,
rtol=1e-03,
atol=1e-06)
np.testing.assert_allclose(th.detach().cpu().numpy(),
jh.data,
rtol=1e-03,
atol=1e-06)
np.testing.assert_allclose(tc.detach().cpu().numpy(),
jc.data,
rtol=1e-03,
atol=1e-06)
def check_equal_1(t_rnn, j_rnn, input, h0):
j_rnn.load_state_dict(t_rnn.state_dict())
t_output, th = t_rnn(torch.from_numpy(input), torch.from_numpy(h0))
j_output, jh = j_rnn(jt.float32(input), jt.float32(h0))
assert np.allclose(t_output.detach().numpy(), j_output.data, rtol=1e-03, atol=1e-06)
assert np.allclose(th.detach().numpy(), jh.data, rtol=1e-03, atol=1e-06)
def check_equal(t_rnn, j_rnn, input, h0, c0):
j_rnn.load_state_dict(t_rnn.state_dict())
t_output, (th, tc) = t_rnn(torch.from_numpy(input),
(torch.from_numpy(h0), torch.from_numpy(c0)))
j_output, (jh, jc) = j_rnn(jt.float32(input),
(jt.float32(h0), jt.float32(c0)))
assert np.allclose(t_output.detach().numpy(), j_output.data, rtol=1e-03, atol=1e-06)
assert np.allclose(th.detach().numpy(), jh.data, rtol=1e-03, atol=1e-06)
assert np.allclose(tc.detach().numpy(), jc.data, rtol=1e-03, atol=1e-06)
def test_var_holder(self):
jt.clean()
expect_error(lambda: jt.matmul(1,1))
expect_error(lambda: jt.matmul([1],[1]))
expect_error(lambda: jt.matmul([[1]],[1]))
self.assertEqual(jt.number_of_lived_vars(), 0)
a = jt.matmul(jt.float32([[3]]), jt.float32([[4]])).data
assert a.shape == (1,1) and a[0,0] == 12
a = np.array([[1, 0], [0, 1]]).astype("float32")
b = np.array([[4, 1], [2, 2]]).astype("float32")
c = np.matmul(a, b)
jtc = jt.matmul(jt.array(a), jt.array(b)).data
assert np.all(jtc == c)
def test_simple_model_train(self):
with jt.flag_scope(trace_py_var=2):
model = Model(input_size=1)
opt = jt.optim.SGD(model.parameters(), 0.1)
batch_size = 10
x = jt.float32(np.random.rand(batch_size, 1))
y = model(x)
opt.step(y**2)
jt.sync_all()
data = jt.dump_trace_data()
jt.clear_trace_data()
# print_stack_tree(data)
for k, v in data["execute_op_info"].items():
for i in v['fused_ops']:
if i not in data["node_data"]:
assert 0, (i, "not found")
for k, v in list(data["node_data"].items()):
if v["attrs"]["name"] == "unname":
assert 0
print(len(data["node_data"]))
with open(f"{jt.flags.cache_path}/simple_model_train.pkl",
"wb") as f:
pickle.dump(data, f)
def test_resnet_infer_with_feature(self):
cat_url = "https://ss1.bdstatic.com/70cFuXSh_Q1YnxGkpoWK1HF6hhy/it/u=3782485413,1118109468&fm=26&gp=0.jpg"
import jittor_utils
cat_path = f"{jt.flags.cache_path}/cat.jpg"
print("download")
jittor_utils.download(cat_url, cat_path)
with open(cat_path, 'rb') as f:
img = Image.open(f).convert('RGB')
img = jt.array(np.array(img))
print(img.shape, img.dtype)
img = ((img.float() - 128) / 255).transpose(2, 0, 1)
with jt.flag_scope(trace_py_var=2, trace_var_data=1):
img = img[None, ...]
resnet18 = resnet.Resnet18(pretrained=True)
x = jt.float32(img)
y = resnet18(x)
y.sync()
data = jt.dump_trace_data()
jt.clear_trace_data()
with open(f"{jt.flags.cache_path}/resnet_with_feature.pkl",
"wb") as f:
pickle.dump(data, f)
for k, v in data["execute_op_info"].items():
for i in v['fused_ops']:
if i not in data["node_data"]:
assert 0, (i, "not found")
def test_stop_fuse2(self):
with jt.profile_scope() as report:
a = jt.float32(0).stop_fuse()
c = jt.float32(0).stop_fuse()
bs = [c]
for i in range(2000):
b = jt.float32(i) * 2 * c
bs.append(b)
a += b
a = a * 2
dbs = jt.grad(a, bs)
jt.sync(dbs + [a])
for a in report[1:]:
assert len(a[0].split("opkey")) < 8
def evaluate(net, epoch, dataloader):
total_acc = 0
total_num = 0
net.eval()
for pts, normals, labels in tqdm(dataloader, desc=f'Epoch {epoch} [Val]'):
pts = jt.float32(pts.numpy())
normals = jt.float32(normals.numpy())
labels = jt.int32(labels.numpy())
output = net(pts, normals)
pred = np.argmax(output.data, axis=1)
acc = np.sum(pred == labels.data)
total_acc += acc
total_num += labels.shape[0]
acc = total_acc / total_num
return acc
def test_stop_fuse(self):
with jt.profile_scope() as report:
a = jt.float32(0).stop_fuse()
c = jt.float32(0)
bs = [c]
for i in range(2000):
b = jt.float32(i) * 2 * c
bs.append(b)
a += b
a = a * 2
dbs = jt.grad(a, bs)
jt.sync(dbs + [a])
for a in report[1:]:
# origin is 50
# after update queue, increase to 102
assert len(a[0].split("opkey")) < 110, len(a[0].split("opkey"))
def test_resnet_train_profile(self):
with jt.profile_scope(trace_py_var=1):
resnet18 = resnet.Resnet18()
opt = jt.optim.SGD(resnet18.parameters(), 0.1)
x = jt.float32(np.random.rand(2, 3, 224, 224))
y = resnet18(x)
opt.step(y**2)
jt.sync_all()
def test_resnet(self):
with jt.flag_scope(trace_py_var=2):
resnet18 = resnet.Resnet18()
x = jt.float32(np.random.rand(2, 3, 224, 224))
y = resnet18(x)
y.sync()
data = jt.dump_trace_data()
jt.clear_trace_data()
def batch_norm(x):
xmean = jt.mean(x, dims=[0, 2, 3], keepdims=1)
x2mean = jt.mean(x * x, dims=[0, 2, 3], keepdims=1)
norm_x = (x - xmean.broadcast_var(x)) / (
jt.sqrt(x2mean - xmean * xmean + jt.float32(1e-5)).broadcast_var(x))
w = jt.make_var([x.shape[1]], init=get_init_var)
b = jt.make_var([x.shape[1]], init=get_init_var)
w = w.broadcast([1, w.shape[0], 1, 1], [0, 2, 3])
b = b.broadcast([1, b.shape[0], 1, 1], [0, 2, 3])
return norm_x * w + b
def test_simple_model(self):
with jt.flag_scope(trace_py_var=2):
model = Model(input_size=1)
batch_size = 10
x = jt.float32(np.random.rand(batch_size, 1))
y = model(x)
y.sync()
data = jt.dump_trace_data()
jt.clear_trace_data()
def test_resnet_train(self):
with jt.flag_scope(trace_py_var=2):
resnet18 = resnet.Resnet18()
opt = jt.optim.SGD(resnet18.parameters(), 0.1)
x = jt.float32(np.random.rand(2, 3, 224, 224))
y = resnet18(x)
opt.step(y**2)
jt.sync_all()
data = jt.dump_trace_data()
jt.clear_trace_data()
def test_simple_model(self):
with jt.flag_scope(trace_py_var=2):
model = Model(input_size=1)
batch_size = 10
x = jt.float32(np.random.rand(batch_size, 1))
y = model(x)
y.sync()
data = jt.dump_trace_data()
jt.clear_trace_data()
with open(f"{jt.flags.cache_path}/simple_model.pkl", "wb") as f:
pickle.dump(data, f)
def test_simple_model_train(self):
with jt.flag_scope(trace_py_var=2):
model = Model(input_size=1)
opt = jt.optim.SGD(model.parameters(), 0.1)
batch_size = 10
x = jt.float32(np.random.rand(batch_size, 1))
y = model(x)
opt.step(y**2)
jt.sync_all()
data = jt.dump_trace_data()
jt.clear_trace_data()
def test_resnet_infer(self):
with jt.flag_scope(trace_py_var=2):
resnet18 = resnet.Resnet18()
x = jt.float32(np.random.rand(2, 3, 224, 224))
y = resnet18(x)
y.sync()
data = jt.dump_trace_data()
jt.clear_trace_data()
with open(f"{jt.flags.cache_path}/resnet.pkl", "wb") as f:
pickle.dump(data, f)
for k, v in data["execute_op_info"].items():
for i in v['fused_ops']:
if i not in data["node_data"]:
assert 0, (i, "not found")
def test_resnet_trainx(self):
with jt.flag_scope(trace_py_var=2):
resnet18 = resnet.Resnet18()
opt = jt.optim.SGD(resnet18.parameters(), 0.1)
x = jt.float32(np.random.rand(2, 3, 224, 224))
y = resnet18(x)
opt.step(y**2)
jt.sync_all()
data = jt.dump_trace_data()
jt.clear_trace_data()
with open(f"{jt.flags.cache_path}/resnet_train.pkl", "wb") as f:
pickle.dump(data, f)
for k, v in data["execute_op_info"].items():
for i in v['fused_ops']:
if i not in data["node_data"]:
assert 0, (i, "not found")
for k, v in data["node_data"].items():
if 'name' not in v["attrs"]:
print(v)
def get_data(n):
for i in range(n):
x = np.random.rand(batch_size, 1)
y = x * x
yield jt.float32(x), jt.float32(y)
import jittor as jt a = jt.float32([1, 2, 3]) print(a) print(a.data)
def relu(x): return jt.maximum(x, jt.float32(0)) def resnet_fake():
def __init__(self, channel):
super(Blur, self).__init__()
self.weight = jt.float32([[1, 2, 1], [2, 4, 2], [1, 2, 1]])
self.weight = self.weight.view(1, 1, 3, 3)
self.weight = self.weight / self.weight.sum()
self.weight.requires_grad = False
def relu(x):
return jt.maximum(x, jt.float32(0))
def get_actions(self, observations):
observations = jt.float32(observations)
if len(observations.shape) == 1:
observations = observations.reshape(1, -1)
return self(observations).data[0]
return self(observations).data
import jittor as jt
import numpy as np
X_train_fpath = sys.argv[1]
Y_train_fpath = sys.argv[2]
X_test_fpath = sys.argv[3]
output_fpath = sys.argv[4]
print('# read training data')
X_train = np.genfromtxt(X_train_fpath, delimiter=',', skip_header=1)
Y_train = np.genfromtxt(Y_train_fpath, delimiter=',', skip_header=1)
X_train = np.delete(X_train, 0, axis=1)
Y_train = np.delete(Y_train, 0, axis=1)
X_train = jt.float32(X_train)
Y_train = jt.float32(Y_train)
def _normalize_column_0_1(X, train=True, specified_column = None, X_min = None, X_max=None):
# The output of the function will make the specified column of the training data
# from 0 to 1
# When processing testing data, we need to normalize by the value
# we used for processing training, so we must save the max value of the
# training data
if train:
if specified_column == None:
specified_column = np.arange(X.shape[1])
length = len(specified_column)
X_max = np.reshape(np.max(X[:, specified_column], 0), (1, length))
X_min = np.reshape(np.min(X[:, specified_column], 0), (1, length))
def test_cuda_pow_grad_nan(self):
a = jt.float32([1, -1, -1000.1])
da = jt.grad(a**2, a)
assert np.isnan(da.data).sum() == 0, da.data
def __init__(self):
self.weight = jt.float32([[1, 2, 1], [2, 4, 2], [1, 2, 1]])
self.weight.requires_grad = False
self.layer1 = nn.Linear(1, 10)
def test_sigmoid_nan(self):
a = jt.float32([1, -1, -1000.1])
da = jt.grad(a.sigmoid(), a)
assert np.isnan(da.data).sum() == 0, da.data
env.seed(seed)
np.random.seed(seed)
jt.seed(seed)
prev_time = time.time()
start_time = prev_time
for episode in range(max_episodes):
# print(np.random.rand(3))
state = env.reset()
done = False
tot_reward = 0
length = 0
entropy_term = jt.float32(0).stop_fuse()
num_step = num_steps // update_step
for steps in range(num_step):
rewards = []
values = []
log_probs = []
for s in range(update_step):
state = jt.float32(state).reshape([1, -1])
value, policy_dist = model(state)
value = value[0, 0]
# print("value: ", value())
action = np.random.choice(num_actions,
p=np.squeeze(policy_dist.data))
# action = 0
