def test_var(self):
a = jt.Var([1, 2, 3])
b = jt.Var([1, 2, 3], "float32")
assert a.dtype == "int32"
assert b.dtype == "float32"
assert (a.numpy() == [1, 2, 3]).all()
assert (b.numpy() == [1, 2, 3]).all()
def inference(self, voxel):
if (len(self.gpu_ids) > 0):
voxel = jt.Var(voxel)
else:
voxel = jt.Var(voxel)
self.netPRS.eval()
with jt.no_grad():
(quat, plane) = self.netPRS(voxel)
return (plane, quat)
def execute(self, voxel, points, cp):
voxel = jt.Var(voxel)
points = jt.Var(points)
cp = jt.Var(cp)
(quat, plane) = self.netPRS(voxel)
(loss_ref, loss_rot) = self.sym_loss(points,
cp,
voxel,
plane=plane,
quat=quat)
(loss_reg_plane, loss_reg_rot) = self.reg_loss(plane=plane,
quat=quat,
weight=self.opt.weight)
return [loss_ref, loss_rot, loss_reg_plane, loss_reg_rot]
def test_pickle(self):
import pickle
a = jt.Var([1, 2, 3, 4])
s = pickle.dumps(a, pickle.HIGHEST_PROTOCOL)
b = pickle.loads(s)
assert isinstance(b, jt.Var)
assert (b.data == [1, 2, 3, 4]).all()
def test_np_array(self):
a = jt.Var([1, 2, 3])
b = np.array(a)
assert (b == [1, 2, 3]).all()