def setUp(self): self.params = Parameters() self.model_pb = Model() self.tensors_pb = self.model_pb.param self.embeddings_pb = self.model_pb.embedding_table_info arr1 = np.random.uniform(size=(3, 4)) tensor1_pb = Tensor(arr1, name="x").to_tensor_pb() arr2 = np.random.uniform(size=(4, 5)) tensor2_pb = Tensor(arr2, name="y").to_tensor_pb() self.tensors_pb.extend([tensor1_pb, tensor2_pb]) self.embedding_table_name = "embedding_1" self.embedding_dim = 10 embedding_pb = EmbeddingTableInfo() embedding_pb.name = self.embedding_table_name embedding_pb.dim = self.embedding_dim embedding_pb.initializer = "uniform" embedding_vectors = np.random.uniform(size=(2, 10)) embedding_indices = np.array([0, 8]) embedding_tensor = Tensor( embedding_vectors, indices=embedding_indices, name=self.embedding_table_name, ) embedding_tensor_pb = embedding_tensor.to_tensor_pb() self.tensors_pb.append(embedding_tensor_pb) self.embeddings_pb.append(embedding_pb)
def test_tensor_data_structure(self): # Test tensor values, without indices arr = np.ndarray(shape=[3, 1, 2, 4], dtype=np.int32) tensor = Tensor(arr) self.assertTrue(np.array_equal(arr, tensor.values)) self.assertTrue(np.array_equal(arr, tensor.to_tf_tensor())) self.assertFalse(tensor.is_indexed_slices()) # Test tensor values, with indices indices = np.array([2, 0, 1]) tensor = Tensor(arr, indices) self.assertTrue(np.array_equal(arr, tensor.values)) self.assertTrue(np.array_equal(indices, tensor.indices)) self.assertTrue(np.array_equal(arr, tensor.to_tf_tensor().values)) self.assertTrue(np.array_equal(indices, tensor.to_tf_tensor().indices)) self.assertTrue(tensor.is_indexed_slices()) # Test round trip # tensor to tensor PB tensor = Tensor(arr, indices, name="test") pb = tensor.to_tensor_pb() self.assertEqual(pb.name, "test") self.assertEqual(pb.dim, [3, 1, 2, 4]) self.assertEqual(pb.dtype, tensor_dtype_pb2.DT_INT32) np.testing.assert_array_equal(pb.indices, indices) # tensor PB to tensor tensor_new = Tensor.from_tensor_pb(pb) self.assertEqual(tensor.name, "test") np.testing.assert_array_equal(tensor_new.values, arr) np.testing.assert_array_equal(tensor_new.indices, indices) # Test Tensor().to_ndarray() values = np.array([[1.0, 2.0], [3.0, 4.0]]) indices = np.array([0, 2]) name = "test" tensor = Tensor(values, indices, name) self.assertRaises(NotImplementedError, tensor.to_ndarray) tensor = Tensor(values, name=name) self.assertTrue(np.allclose(values, tensor.to_ndarray()))