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()