def save_ragged_vals_to_dataset(vals_list, output_path, concat_all=True): # print(vals_list) def data_gen(): if concat_all: yield tf.ragged.constant(vals_list, dtype=tf.int32) else: for i, vals in enumerate(vals_list): if i % 10000 == 0: print(i) yield tf.ragged.constant([vals], dtype=tf.int32) dataset = tf.data.Dataset.from_generator( data_gen, output_signature=tf.RaggedTensorSpec(ragged_rank=1, dtype=tf.int32)) # for v in dataset: # print(v) print('saving to', output_path) tf.data.experimental.save(dataset, output_path, shard_func=lambda x: np.int64(0)) print('saved')
def _filter_poses(self, boxes, poses3d, poses2d): poses3d_mean = tf.reduce_mean(poses3d, axis=-3) poses2d_mean = tf.reduce_mean(poses2d, axis=-3) plausible_mask_flat = tf.logical_and( tf.logical_and( is_pose_plausible(poses3d_mean.flat_values), are_augmentation_results_consistent(poses3d.flat_values)), is_pose_consistent_with_box(poses2d_mean.flat_values, boxes.flat_values)) plausible_mask = tf.RaggedTensor.from_row_lengths( plausible_mask_flat, boxes.row_lengths()) # Apply pose similarity-based non-maximum suppression to reduce duplicates nms_indices = tf.map_fn( fn=lambda args: pose_non_max_suppression(*args), elems=(poses3d_mean, boxes[..., 4], plausible_mask), fn_output_signature=tf.RaggedTensorSpec(shape=(None,), dtype=tf.int32)) return nms_indices
def _truncate_row_lengths(ragged_tensor: tf.RaggedTensor, new_lengths: tf.Tensor) -> tf.RaggedTensor: """Truncates the rows of `ragged_tensor` to the given row lengths.""" new_lengths = tf.broadcast_to(new_lengths, ragged_tensor.bounding_shape()[0:1]) def fn(x): row, new_length = x return row[0:new_length] fn_dtype = tf.RaggedTensorSpec(dtype=ragged_tensor.dtype, ragged_rank=ragged_tensor.ragged_rank - 1) result = tf.map_fn(fn, (ragged_tensor, new_lengths), dtype=fn_dtype) # Work around broken shape propagation: without this, result has unknown rank. flat_values_shape = [None] * ragged_tensor.flat_values.shape.rank result = result.with_flat_values( tf.ensure_shape(result.flat_values, flat_values_shape)) return result
def get_input_info_dict(self, signature=None, tags=None): if signature == "ragged" and tags == set(["special"]): result = { "x": tensor_info.ParsedTensorInfo.from_type_spec( type_spec=tf.RaggedTensorSpec( shape=[None, None, None, 3], dtype=tf.float32, ragged_rank=2)), } else: result = { "x": tensor_info.ParsedTensorInfo( tf.float32, tf.TensorShape([None]), is_sparse=(signature == "sparse" and tags == set(["special"]))), } if tags == set(["special"]) and signature == "extra": result["y"] = result["x"] return result
def get_traing_set(self, batch_size, seed): training_set = tf.data.experimental.load( "./train", tf.RaggedTensorSpec(tf.TensorShape([3, None]), tf.int32, 1, tf.int64)) with open('./utils/tid_2_aid') as file: tid_2_aid = tf.constant(json.load(file)) file.close() self.tid_2_aid = tf.lookup.StaticHashTable( tf.lookup.KeyValueTensorInitializer(tid_2_aid[:, 0], tid_2_aid[:, 1]), default_value=-1) del tid_2_aid tf.random.set_seed(seed) np.random.seed(seed) return training_set.map(lambda x: self.corrupt(x)).shuffle( 1000, seed, True).apply( tf.data.experimental.dense_to_ragged_batch( batch_size, drop_remainder=True))
def __init__(self, model_path, antialias_factor=4): super().__init__() self.antialias_factor = antialias_factor self.crop_model = tf.saved_model.load(model_path) self.crop_side = 256 self.joint_names = self.crop_model.joint_names self.joint_edges = self.crop_model.joint_edges joint_names = [b.decode('utf8') for b in self.joint_names.numpy()] self.joint_info = data.datasets3d.JointInfo(joint_names, self.joint_edges.numpy()) self.__call__.get_concrete_function( tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8), tf.TensorSpec(shape=(None, 3, 3), dtype=tf.float32), tf.RaggedTensorSpec(shape=(None, None, 4), ragged_rank=1, dtype=tf.float32), tf.TensorSpec(shape=(), dtype=tf.int32), tf.TensorSpec(shape=(), dtype=tf.int32)) self.__call__.get_concrete_function( tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8), tf.TensorSpec(shape=(3, 3), dtype=tf.float32), tf.TensorSpec(shape=(None, 4), dtype=tf.float32), tf.TensorSpec(shape=(), dtype=tf.int32), tf.TensorSpec(shape=(), dtype=tf.int32))
def hash_matrix(self): ''' Creates hash_matrix with the weight indices ''' indx = tf.convert_to_tensor([[ xxhash.xxh32("{}_{}_{}".format(i, j, self.units), self.hash_seed).intdigest() % self.n_weights for j in range(self.units) ] for i in range(self.input_dim)], dtype=tf.int32) spec = tf.RaggedTensorSpec(dtype=tf.dtypes.int64, ragged_rank=0, row_splits_dtype=tf.dtypes.int64) cond = tf.stack([indx == i for i in range(self.n_weights)]) rag = tf.stack([ tf.map_fn(fn=lambda t: tf.squeeze(tf.where(t), axis=1), elems=c, fn_output_signature=spec, back_prop=False) for c in tf.transpose(cond, [2, 0, 1]) ]) self.indices = (rag + 1).to_tensor()
def testInputSpecsToTensorRepresentations(self): tensor_representations = model_util.input_specs_to_tensor_representations( { 'input_1': tf.TensorSpec(shape=(None, 2), dtype=tf.int64), 'input_2': tf.SparseTensorSpec(shape=(None, 1), dtype=tf.float32), 'input_3': tf.RaggedTensorSpec(shape=(None, None), dtype=tf.float32), }) dense_tensor_representation = text_format.Parse( """ dense_tensor { column_name: "input_1" shape { dim { size: 2 } } } """, schema_pb2.TensorRepresentation()) sparse_tensor_representation = text_format.Parse( """ varlen_sparse_tensor { column_name: "input_2" } """, schema_pb2.TensorRepresentation()) ragged_tensor_representation = text_format.Parse( """ ragged_tensor { feature_path { step: "input_3" } } """, schema_pb2.TensorRepresentation()) self.assertEqual( { 'input_1': dense_tensor_representation, 'input_2': sparse_tensor_representation, 'input_3': ragged_tensor_representation }, tensor_representations)
def convert(self, tensor: TensorAlike) -> List[pa.Array]: """Converts the given TensorAlike to pa.Arrays after validating its spec.""" if tf.__version__ < "2": if isinstance(tensor, np.ndarray): actual_spec = tf.TensorSpec(tensor.shape, tf.dtypes.as_dtype(tensor.dtype)) elif isinstance(tensor, tf.compat.v1.SparseTensorValue): actual_spec = tf.SparseTensorSpec(tensor.dense_shape, tensor.values.dtype) elif isinstance(tensor, tf.compat.v1.ragged.RaggedTensorValue): actual_spec = tf.RaggedTensorSpec( tensor.shape, tensor.values.dtype, row_splits_dtype=tensor.row_splits.dtype) else: raise TypeError("Only ndarrays, SparseTensorValues and " "RaggedTensorValues are supported with TF 1.x, " "got {}".format(type(tensor))) else: actual_spec = tf.type_spec_from_value(tensor) if not self._type_spec.is_compatible_with(actual_spec): raise TypeError("Expected {} but got {}".format(self._type_spec, actual_spec)) return self._convert_internal(tensor)
def testInferTensorSpecs(self): sparse_value = types.SparseTensorValue( values=np.array([0.5, -1., 0.5, -1.], dtype=np.float32), indices=np.array([[0, 3, 1], [0, 20, 0], [1, 3, 1], [1, 20, 0]]), dense_shape=np.array([2, 100, 3])) ragged_value = types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 7, 1, 8, 8, 2, 1], dtype=np.float32), nested_row_splits=[ np.array([0, 3, 6]), np.array([0, 2, 3, 4, 5, 5, 8]), np.array([0, 2, 3, 3, 6, 9, 10, 11, 13]) ]) tensor_values = { 'features': { 'feature_1': np.array([1, 2, 3], dtype=np.float32), 'feature_2': sparse_value, 'feature_3': ragged_value, }, 'labels': np.array([1], dtype=np.float32), } expected_specs = { 'features': { 'feature_1': tf.TensorSpec([None], dtype=tf.float32), 'feature_2': tf.SparseTensorSpec(shape=[None, 100, 3], dtype=tf.float32), 'feature_3': tf.RaggedTensorSpec(shape=[None, None, None, None], dtype=tf.float32) }, 'labels': tf.TensorSpec([None], dtype=tf.float32) } got_specs = util.infer_tensor_specs( util.to_tensorflow_tensors(tensor_values)) self.assertDictEqual(expected_specs, got_specs)
def test_ragged_roundtrip(self, exported_in_tf1): if not hasattr(meta_graph_pb2.TensorInfo, 'CompositeTensor'): self.skipTest('This version of TensorFlow does not support ' 'CompositeTenors in TensorInfo.') input_specs = { 'input': tf.RaggedTensorSpec( shape=[None, None], dtype=tf.float32, ragged_rank=1, row_splits_dtype=tf.int64) } def preprocessing_fn(inputs): return {'output': inputs['input'] / 2.0} export_path = _create_test_saved_model( exported_in_tf1, input_specs, preprocessing_fn, base_dir=self.get_temp_dir()) splits = np.array([0, 2, 3], dtype=np.int64) values = np.array([1.0, 2.0, 4.0], dtype=np.float32) input_ragged = tf.RaggedTensor.from_row_splits(values, splits) # Using a computed input gives confidence that the graphs are fused inputs = {'input': input_ragged * 10} saved_model_loader = saved_transform_io_v2.SavedModelLoader(export_path) outputs = saved_model_loader.apply_transform_model(inputs) result = outputs['output'] self.assertIsInstance(result, tf.RaggedTensor) # indices and shape unchanged; values multipled by 10 and divided by 2 self.assertAllEqual(splits, result.row_splits) self.assertEqual([5.0, 10.0, 20.0], result.values.numpy().tolist())
class TensorToArrowTest(tf.test.TestCase, parameterized.TestCase): def _assert_tensor_alike_equal(self, left, right): self.assertIsInstance(left, type(right)) if isinstance(left, tf.SparseTensor): self.assertAllEqual(left.values, right.values) self.assertAllEqual(left.indices, right.indices) self.assertAllEqual(left.dense_shape, right.dense_shape) else: self.assertAllEqual(left, right) @parameterized.named_parameters(*_CONVERT_TEST_CASES) def test_convert(self, type_specs, expected_schema, expected_tensor_representations, tensor_input, expected_record_batch): converter = tensor_to_arrow.TensorsToRecordBatchConverter(type_specs) expected_schema = pa.schema( [pa.field(n, t) for n, t in sorted(expected_schema.items())]) self.assertTrue(converter.arrow_schema().equals(expected_schema), "actual: {}".format(converter.arrow_schema())) canonical_expected_tensor_representations = {} for n, r in expected_tensor_representations.items(): if not isinstance(r, schema_pb2.TensorRepresentation): r = text_format.Parse(r, schema_pb2.TensorRepresentation()) canonical_expected_tensor_representations[n] = r self.assertEqual(canonical_expected_tensor_representations, converter.tensor_representations()) rb = converter.convert(tensor_input) self.assertTrue( rb.equals( pa.record_batch( [arr for _, arr in sorted(expected_record_batch.items())], schema=expected_schema))) # Test that TensorAdapter(TensorsToRecordBatchConverter()) is identity. adapter = tensor_adapter.TensorAdapter( tensor_adapter.TensorAdapterConfig( arrow_schema=converter.arrow_schema(), tensor_representations=converter.tensor_representations())) adapter_output = adapter.ToBatchTensors(rb, produce_eager_tensors=True) self.assertEqual(adapter_output.keys(), tensor_input.keys()) for k in adapter_output.keys(): self._assert_tensor_alike_equal(adapter_output[k], tensor_input[k]) def test_unable_to_handle(self): with self.assertRaisesRegex(ValueError, "No handler found"): tensor_to_arrow.TensorsToRecordBatchConverter( {"sp": tf.SparseTensorSpec([None, None, None], tf.int32)}) with self.assertRaisesRegex(ValueError, "No handler found"): tensor_to_arrow.TensorsToRecordBatchConverter( {"sp": tf.SparseTensorSpec([None, None], tf.bool)}) def test_incompatible_type_spec(self): converter = tensor_to_arrow.TensorsToRecordBatchConverter( {"sp": tf.SparseTensorSpec([None, None], tf.int32)}) with self.assertRaisesRegex(TypeError, "Expected SparseTensorSpec"): converter.convert({ "sp": tf.SparseTensor(indices=[[0, 1]], values=tf.constant([0], dtype=tf.int64), dense_shape=[4, 1]) }) @parameterized.named_parameters(*[ dict(testcase_name="bool_value_type", spec=tf.RaggedTensorSpec(shape=[2, None, None], dtype=tf.bool, ragged_rank=2, row_splits_dtype=tf.int64)), dict(testcase_name="2d_leaf_value", spec=tf.RaggedTensorSpec(shape=[2, None, None], dtype=tf.int32, ragged_rank=1, row_splits_dtype=tf.int64)), dict(testcase_name="ragged_rank_less_than_one", spec=tf.RaggedTensorSpec(shape=[2], dtype=tf.int32, ragged_rank=0, row_splits_dtype=tf.int64)), ]) def test_unable_to_handle_ragged(self, spec): with self.assertRaisesRegex(ValueError, "No handler found"): tensor_to_arrow.TensorsToRecordBatchConverter({"rt": spec})
tf.SparseTensor(values=[b"aa", b"bb"], indices=[[2, 0], [2, 1]], dense_shape=[4, 2]) }, expected_record_batch={ "sp1": pa.array([[1], [], [2], []], type=pa.large_list(pa.int32())), "sp2": pa.array([[], [], [b"aa", b"bb"], []], type=pa.large_list(pa.large_binary())) }), dict(testcase_name="ragged_tensors", type_specs={ "sp1": tf.RaggedTensorSpec(tf.TensorShape([3, None]), tf.int64, ragged_rank=1, row_splits_dtype=tf.int64), "sp2": tf.RaggedTensorSpec(tf.TensorShape([3, None]), tf.string, ragged_rank=1, row_splits_dtype=tf.int64), }, expected_schema={ "sp1": pa.large_list(pa.int64()), "sp2": pa.large_list(pa.large_binary()), }, expected_tensor_representations={ "sp1": """ragged_tensor { feature_path {
class Pose3dEstimator(tf.Module): def __init__(self): super().__init__() # Note that only the Trackable resource attributes such as Variables and Models will be # retained when saving to SavedModel self.crop_model = tf.saved_model.load(FLAGS.input_model_path) self.crop_side = FLAGS.crop_side self.joint_names = self.crop_model.joint_names self.joint_edges = self.crop_model.joint_edges joint_names = [b.decode('utf8') for b in self.joint_names.numpy()] self.joint_info = data.datasets3d.JointInfo(joint_names, self.joint_edges.numpy()) self.detector = tf.saved_model.load(FLAGS.detector_path) if FLAGS.detector_path else None if len(joint_names) == 122: skeleton_infos = util.load_pickle('./saved_model_export/skeleton_types.pkl') self.per_skeleton_indices = { k: tf.Variable(v['indices'], dtype=tf.int32, trainable=False) for k, v in skeleton_infos.items()} self.per_skeleton_joint_names = { k: tf.Variable(v['names'], dtype=tf.string, trainable=False) for k, v in skeleton_infos.items()} self.per_skeleton_joint_edges = { k: tf.Variable(v['edges'], dtype=tf.int32, trainable=False) for k, v in skeleton_infos.items()} self.per_skeleton_indices[''] = tf.range(122, dtype=tf.int32) self.per_skeleton_joint_names[''] = self.joint_names self.per_skeleton_joint_edges[''] = self.joint_edges @tf.function(input_signature=[ tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8), # images tf.TensorSpec(shape=(None, 3, 3), dtype=tf.float32), # intrinsic_matrix tf.TensorSpec(shape=(None, 5), dtype=tf.float32), # distortion_coeffs tf.TensorSpec(shape=(None, 4, 4), dtype=tf.float32), # extrinsic_matrix tf.TensorSpec(shape=(3,), dtype=tf.float32), # world_up_vector tf.TensorSpec(shape=(), dtype=tf.float32), # default_fov_degrees tf.TensorSpec(shape=(), dtype=tf.int32), # internal_batch_size tf.TensorSpec(shape=(), dtype=tf.int32), # antialias_factor tf.TensorSpec(shape=(), dtype=tf.int32), # num_aug tf.TensorSpec(shape=(), dtype=tf.bool), # average_aug tf.TensorSpec(shape=(), dtype=tf.string), # skeleton tf.TensorSpec(shape=(), dtype=tf.float32), # detector_threshold tf.TensorSpec(shape=(), dtype=tf.float32), # detector_nms_iou_threshold tf.TensorSpec(shape=(), dtype=tf.int32), # max_detections tf.TensorSpec(shape=(), dtype=tf.bool), # detector_flip_aug tf.TensorSpec(shape=(), dtype=tf.bool)]) # suppress_implausible_poses def detect_poses_batched( self, images, intrinsic_matrix=(UNKNOWN_INTRINSIC_MATRIX,), distortion_coeffs=(DEFAULT_DISTORTION,), extrinsic_matrix=(DEFAULT_EXTRINSIC_MATRIX,), world_up_vector=DEFAULT_WORLD_UP, default_fov_degrees=55, internal_batch_size=64, antialias_factor=1, num_aug=5, average_aug=True, skeleton='', detector_threshold=0.3, detector_nms_iou_threshold=0.7, max_detections=-1, detector_flip_aug=False, suppress_implausible_poses=True): boxes = self._get_boxes( images, detector_flip_aug, detector_nms_iou_threshold, detector_threshold, max_detections) return self._estimate_poses_batched( images, boxes, intrinsic_matrix, distortion_coeffs, extrinsic_matrix, world_up_vector, default_fov_degrees, internal_batch_size, antialias_factor, num_aug, average_aug, skeleton, suppress_implausible_poses) @tf.function(input_signature=[ tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8), # images tf.RaggedTensorSpec(shape=(None, None, 4), ragged_rank=1, dtype=tf.float32), # boxes tf.TensorSpec(shape=(None, 3, 3), dtype=tf.float32), # intrinsic_matrix tf.TensorSpec(shape=(None, 5), dtype=tf.float32), # distortion_coeffs tf.TensorSpec(shape=(None, 4, 4), dtype=tf.float32), # extrinsic_matrix tf.TensorSpec(shape=(3,), dtype=tf.float32), # world_up_vector tf.TensorSpec(shape=(), dtype=tf.float32), # default_fov_degrees tf.TensorSpec(shape=(), dtype=tf.int32), # internal_batch_size tf.TensorSpec(shape=(), dtype=tf.int32), # antialias_factor tf.TensorSpec(shape=(), dtype=tf.int32), # num_aug tf.TensorSpec(shape=(), dtype=tf.bool), # average_aug tf.TensorSpec(shape=(), dtype=tf.string), # skeleton ]) def estimate_poses_batched( self, images, boxes, intrinsic_matrix=(UNKNOWN_INTRINSIC_MATRIX,), distortion_coeffs=(DEFAULT_DISTORTION,), extrinsic_matrix=(DEFAULT_EXTRINSIC_MATRIX,), world_up_vector=DEFAULT_WORLD_UP, default_fov_degrees=55, internal_batch_size=64, antialias_factor=1, num_aug=5, average_aug=True, skeleton=''): boxes = tf.concat([boxes, tf.ones_like(boxes[..., :1])], axis=-1) pred = self._estimate_poses_batched( images, boxes, intrinsic_matrix, distortion_coeffs, extrinsic_matrix, world_up_vector, default_fov_degrees, internal_batch_size, antialias_factor, num_aug, average_aug, skeleton, suppress_implausible_poses=False) del pred['boxes'] return pred def _estimate_poses_batched( self, images, boxes, intrinsic_matrix, distortion_coeffs, extrinsic_matrix, world_up_vector, default_fov_degrees, internal_batch_size, antialias_factor, num_aug, average_aug, skeleton, suppress_implausible_poses): # Special case when zero boxes are provided or found # (i.e., all images images without person detections) # This must be explicitly handled, else the shapes don't work out automatically # for the TensorArray in _predict_in_batches. if tf.size(boxes) == 0: return self._predict_empty(images, num_aug, average_aug) n_images = tf.shape(images)[0] # If one intrinsic matrix is given, repeat it for all images if tf.shape(intrinsic_matrix)[0] == 1: # If intrinsic_matrix is not given, fill it in based on field of view if tf.reduce_all(intrinsic_matrix == -1): intrinsic_matrix = intrinsic_matrix_from_field_of_view( default_fov_degrees, tf.shape(images)[1:3]) intrinsic_matrix = tf.repeat(intrinsic_matrix, n_images, axis=0) # If one distortion coeff/extrinsic matrix is given, repeat it for all images if tf.shape(distortion_coeffs)[0] == 1: distortion_coeffs = tf.repeat(distortion_coeffs, n_images, axis=0) if tf.shape(extrinsic_matrix)[0] == 1: extrinsic_matrix = tf.repeat(extrinsic_matrix, n_images, axis=0) # Now repeat these camera params for each box n_box_per_image = boxes.row_lengths() intrinsic_matrix = tf.repeat(intrinsic_matrix, n_box_per_image, axis=0) distortion_coeffs = tf.repeat(distortion_coeffs, n_box_per_image, axis=0) # Up-vector in camera-space camspace_up = tf.einsum('c,bCc->bC', world_up_vector, extrinsic_matrix[..., :3, :3]) camspace_up = tf.repeat(camspace_up, n_box_per_image, axis=0) # Set up the test-time augmentation parameters aug_gammas = tf.cast(tf.linspace(0.6, 1.0, num_aug), tf.float32) aug_angle_range = np.float32(np.deg2rad(FLAGS.rot_aug)) if FLAGS.rot_aug_linspace_noend: aug_angles = linspace_noend(-aug_angle_range, aug_angle_range, num_aug) else: aug_angles = tf.linspace(-aug_angle_range, aug_angle_range, num_aug) aug_scales = tf.concat([ linspace_noend(0.8, 1.0, num_aug // 2), tf.linspace(1.0, 1.1, num_aug - num_aug // 2)], axis=0) aug_should_flip = (tf.range(num_aug) - num_aug // 2) % 2 != 0 aug_flipmat = tf.constant([[-1, 0, 0], [0, 1, 0], [0, 0, 1]], np.float32) aug_maybe_flipmat = tf.where( aug_should_flip[:, np.newaxis, np.newaxis], aug_flipmat, tf.eye(3)) aug_rotmat = rotation_mat_zaxis(-aug_angles) aug_rotflipmat = aug_maybe_flipmat @ aug_rotmat # crops_flat, poses3dcam_flat = self._predict_in_batches( poses3d_flat = self._predict_in_batches( images, intrinsic_matrix, distortion_coeffs, camspace_up, boxes, internal_batch_size, aug_should_flip, aug_rotflipmat, aug_gammas, aug_scales, antialias_factor) # Project the 3D poses to get the 2D poses poses2d_flat_normalized = to_homogeneous( distort_points(project(poses3d_flat), distortion_coeffs)) poses2d_flat = tf.einsum('bank,bjk->banj', poses2d_flat_normalized, intrinsic_matrix[..., :2, :]) poses2d_flat = tf.ensure_shape(poses2d_flat, [None, None, self.joint_info.n_joints, 2]) # Arrange the results back into ragged tensors poses3d = tf.RaggedTensor.from_row_lengths(poses3d_flat, n_box_per_image) poses2d = tf.RaggedTensor.from_row_lengths(poses2d_flat, n_box_per_image) # crops = tf.RaggedTensor.from_row_lengths(crops_flat, n_box_per_image) if suppress_implausible_poses: # Filter the resulting poses for individual plausibility to reduce false positives selected_indices = self._filter_poses(boxes, poses3d, poses2d) boxes, poses3d, poses2d = [ tf.gather(x, selected_indices, batch_dims=1) for x in [boxes, poses3d, poses2d]] # crops = tf.gather(crops, selected_indices, batch_dims=1) # Convert to world coordinates extrinsic_matrix = tf.repeat(tf.linalg.inv(extrinsic_matrix), poses3d.row_lengths(), axis=0) poses3d = tf.RaggedTensor.from_row_lengths( tf.einsum( 'bank,bjk->banj', to_homogeneous(poses3d.flat_values), extrinsic_matrix[..., :3, :]), poses3d.row_lengths()) if skeleton != '': poses3d = self._get_skeleton(poses3d, skeleton) poses2d = self._get_skeleton(poses2d, skeleton) if average_aug: poses3d = tf.reduce_mean(poses3d, axis=-3) poses2d = tf.reduce_mean(poses2d, axis=-3) result = dict(boxes=boxes, poses3d=poses3d, poses2d=poses2d) # result['crops'] = crops return result def _get_boxes( self, images, detector_flip_aug, detector_nms_iou_threshold, detector_threshold, max_detections): if self.detector is None: n_images = tf.shape(images)[0] boxes = tf.RaggedTensor.from_row_lengths( tf.zeros(shape=(0, 5)), tf.zeros(n_images, tf.int64)) else: boxes = self.detector.predict_multi_image( images, detector_threshold, detector_nms_iou_threshold, detector_flip_aug, detector_flip_aug and FLAGS.detector_flip_vertical_too) if max_detections > -1 and not tf.size(boxes) == 0: topk_indices = topk_indices_ragged(boxes[..., 4], max_detections) boxes = tf.gather(boxes, topk_indices, axis=1, batch_dims=1) return boxes def _predict_in_batches( self, images, intrinsic_matrix, distortion_coeffs, camspace_up, boxes, internal_batch_size, aug_should_flip, aug_rotflipmat, aug_gammas, aug_scales, antialias_factor): num_aug = tf.shape(aug_gammas)[0] boxes_per_batch = internal_batch_size // num_aug boxes_flat = boxes.flat_values image_id_per_box = boxes.value_rowids() # Gamma decoding for correct image rescaling later on images = (tf.cast(images, tf.float32) / np.float32(255)) ** 2.2 if boxes_per_batch == 0: # Run all as a single batch return self._predict_single_batch( images, intrinsic_matrix, distortion_coeffs, camspace_up, boxes_flat, image_id_per_box, aug_rotflipmat, aug_should_flip, aug_scales, aug_gammas, antialias_factor) else: # Chunk the image crops into batches and predict them one by one n_total_boxes = tf.shape(boxes_flat)[0] n_batches = tf.cast(tf.math.ceil(n_total_boxes / boxes_per_batch), tf.int32) poses3d_batches = tf.TensorArray( tf.float32, size=n_batches, element_shape=(None, None, self.joint_info.n_joints, 3), infer_shape=False) # crop_batches = tf.TensorArray( # tf.float32, size=n_batches, # element_shape=(None, None, self.crop_side, self.crop_side, 3), # infer_shape=False) for i in tf.range(n_batches): batch_slice = slice(i * boxes_per_batch, (i + 1) * boxes_per_batch) # crops, poses3d = self._predict_single_batch( poses3d = self._predict_single_batch( images, intrinsic_matrix[batch_slice], distortion_coeffs[batch_slice], camspace_up[batch_slice], boxes_flat[batch_slice], image_id_per_box[batch_slice], aug_rotflipmat, aug_should_flip, aug_scales, aug_gammas, antialias_factor) poses3d_batches = poses3d_batches.write(i, poses3d) # crop_batches = crop_batches.write(i, crops) # return crop_batches.concat(), poses3d_batches.concat() return poses3d_batches.concat() def _predict_single_batch( self, images, intrinsic_matrix, distortion_coeffs, camspace_up, boxes, image_ids, aug_rotflipmat, aug_should_flip, aug_scales, aug_gammas, antialias_factor): # Get crops and info about the transformation used to create them # Each has shape [num_aug, n_boxes, ...] crops, new_intrinsic_matrix, R = self._get_crops( images, intrinsic_matrix, distortion_coeffs, camspace_up, boxes, image_ids, aug_rotflipmat, aug_scales, aug_gammas, antialias_factor) # Flatten each and predict the pose with the crop model new_intrinsic_matrix_flat = tf.reshape(new_intrinsic_matrix, (-1, 3, 3)) crops_flat = tf.reshape(crops, (-1, self.crop_side, self.crop_side, 3)) poses_flat = self.crop_model.predict_multi( tf.cast(crops_flat, tf.float16), new_intrinsic_matrix_flat) # Unflatten the result num_aug = tf.shape(aug_should_flip)[0] poses = tf.reshape(poses_flat, [num_aug, -1, self.joint_info.n_joints, 3]) # Reorder the joints for mirrored predictions (e.g., swap left and right wrist) left_right_swapped_poses = tf.gather(poses, self.joint_info.mirror_mapping, axis=-2) poses = tf.where( tf.reshape(aug_should_flip, [-1, 1, 1, 1]), left_right_swapped_poses, poses) # Transform the predictions back into the original camera space # We need to multiply by the inverse of R, but since we are using row vectors in `poses` # the inverse and transpose cancel out, leaving just R. poses_orig_camspace = poses @ R # Transpose to [n_boxes, num_aug, ...] # return (tf.transpose(crops, [1, 0, 2, 3, 4]), # tf.transpose(poses_orig_camspace, [1, 0, 2, 3])) return tf.transpose(poses_orig_camspace, [1, 0, 2, 3]) def _get_crops( self, images, intrinsic_matrix, distortion_coeffs, camspace_up, boxes, image_ids, aug_rotflipmat, aug_scales, aug_gammas, antialias_factor): R_noaug, box_aug_scales = self._get_new_rotation_and_scale( intrinsic_matrix, distortion_coeffs, camspace_up, boxes) # How much we need to scale overall, taking scale augmentation into account # From here on, we introduce the dimension of augmentations crop_scales = aug_scales[:, tf.newaxis] * box_aug_scales[tf.newaxis, :] # Build the new intrinsic matrix n_box = tf.shape(boxes)[0] num_aug = tf.shape(aug_gammas)[0] new_intrinsic_matrix = tf.concat([ tf.concat([ # Top-left of original intrinsic matrix gets scaled intrinsic_matrix[tf.newaxis, :, :2, :2] * crop_scales[:, :, tf.newaxis, tf.newaxis], # Principal point is the middle of the new image size tf.fill((num_aug, n_box, 2, 1), self.crop_side / 2)], axis=3), tf.concat([ # [0, 0, 1] as the last row of the intrinsic matrix: tf.zeros((num_aug, n_box, 1, 2), tf.float32), tf.ones((num_aug, n_box, 1, 1), tf.float32)], axis=3)], axis=2) R = aug_rotflipmat[:, tf.newaxis] @ R_noaug new_invprojmat = tf.linalg.inv(new_intrinsic_matrix @ R) # If we perform antialiasing through output scaling, we render a larger image first and then # shrink it. So we scale the homography first. if antialias_factor > 1: scaling_mat = corner_aligned_scale_mat(1 / tf.cast(antialias_factor, tf.float32)) new_invprojmat = new_invprojmat @ scaling_mat crops = warp_images( images, intrinsic_matrix=tf.tile(intrinsic_matrix, [num_aug, 1, 1]), new_invprojmat=tf.reshape(new_invprojmat, [-1, 3, 3]), distortion_coeffs=tf.tile(distortion_coeffs, [num_aug, 1]), crop_scales=tf.reshape(crop_scales, [-1]) * tf.cast(antialias_factor, tf.float32), output_shape=(self.crop_side * antialias_factor, self.crop_side * antialias_factor), image_ids=tf.tile(image_ids, [num_aug])) # Downscale the result if we do antialiasing through output scaling if antialias_factor == 2: crops = tf.nn.avg_pool2d(crops, 2, 2, padding='VALID') elif antialias_factor == 4: crops = tf.nn.avg_pool2d(crops, 4, 4, padding='VALID') elif antialias_factor > 4: crops = tf.image.resize( crops, (self.crop_side, self.crop_side), tf.image.ResizeMethod.AREA) crops = tf.reshape(crops, [num_aug, n_box, self.crop_side, self.crop_side, 3]) # The division by 2.2 cancels the original gamma decoding from earlier crops **= tf.reshape(aug_gammas / 2.2, [-1, 1, 1, 1, 1]) return crops, new_intrinsic_matrix, R def _get_new_rotation_and_scale(self, intrinsic_matrix, distortion_coeffs, camspace_up, boxes): # Transform five points on each box: the center and the four side midpoints x, y, w, h = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] boxpoints_homog = to_homogeneous(tf.stack([ tf.stack([x + w / 2, y + h / 2], axis=1), tf.stack([x + w / 2, y], axis=1), tf.stack([x + w, y + h / 2], axis=1), tf.stack([x + w / 2, y + h], axis=1), tf.stack([x, y + h / 2], axis=1)], axis=1)) boxpoints_camspace = tf.einsum( 'bpc,bCc->bpC', boxpoints_homog, tf.linalg.inv(intrinsic_matrix)) boxpoints_camspace = to_homogeneous( undistort_points(boxpoints_camspace[..., :2], distortion_coeffs)) # Create a rotation matrix that will put the box center to the principal point # and apply the augmentation rotation and flip, to get the new coordinate frame box_center_camspace = boxpoints_camspace[:, 0] R_noaug = get_new_rotation_matrix(forward_vector=box_center_camspace, up_vector=camspace_up) # Transform the side midpoints of the box to the new coordinate frame sidepoints_camspace = boxpoints_camspace[:, 1:5] sidepoints_new = project(tf.einsum( 'bpc,bCc->bpC', sidepoints_camspace, intrinsic_matrix @ R_noaug)) # Measure the size of the reprojected boxes vertical_size = tf.linalg.norm(sidepoints_new[:, 0] - sidepoints_new[:, 2], axis=-1) horiz_size = tf.linalg.norm(sidepoints_new[:, 1] - sidepoints_new[:, 3], axis=-1) box_size_new = tf.maximum(vertical_size, horiz_size) # How much we need to scale (zoom) to have the boxes fill out the final crop box_aug_scales = self.crop_side / box_size_new return R_noaug, box_aug_scales def _predict_empty(self, image, num_aug, average_aug): if average_aug: poses3d = tf.zeros(shape=(0, self.joint_info.n_joints, 3)) poses2d = tf.zeros(shape=(0, self.joint_info.n_joints, 2)) else: poses3d = tf.zeros(shape=(0, num_aug, self.joint_info.n_joints, 3)) poses2d = tf.zeros(shape=(0, num_aug, self.joint_info.n_joints, 2)) n_images = tf.shape(image)[0] poses3d = tf.RaggedTensor.from_row_lengths(poses3d, tf.zeros(n_images, tf.int64)) poses2d = tf.RaggedTensor.from_row_lengths(poses2d, tf.zeros(n_images, tf.int64)) boxes = tf.zeros(shape=(0, 5)) boxes = tf.RaggedTensor.from_row_lengths(boxes, tf.zeros(n_images, tf.int64)) result = dict(boxes=boxes, poses3d=poses3d, poses2d=poses2d) # crops = tf.zeros(shape=(0, num_aug, self.crop_side, self.crop_side, 3)) # crops = tf.RaggedTensor.from_row_lengths(crops, tf.zeros(n_images, tf.int64)) # result['crops'] = crops return result def _filter_poses(self, boxes, poses3d, poses2d): poses3d_mean = tf.reduce_mean(poses3d, axis=-3) poses2d_mean = tf.reduce_mean(poses2d, axis=-3) plausible_mask_flat = tf.logical_and( tf.logical_and( is_pose_plausible(poses3d_mean.flat_values), are_augmentation_results_consistent(poses3d.flat_values)), is_pose_consistent_with_box(poses2d_mean.flat_values, boxes.flat_values)) plausible_mask = tf.RaggedTensor.from_row_lengths( plausible_mask_flat, boxes.row_lengths()) # Apply pose similarity-based non-maximum suppression to reduce duplicates nms_indices = tf.map_fn( fn=lambda args: pose_non_max_suppression(*args), elems=(poses3d_mean, boxes[..., 4], plausible_mask), fn_output_signature=tf.RaggedTensorSpec(shape=(None,), dtype=tf.int32)) return nms_indices def _get_skeleton(self, poses, skeleton): # We must list all possibilities since we can't address the Python dictionary # `self.per_skeleton_indices` with the tf.Tensor `skeleton`. if skeleton == b'smpl_24': indices = self.per_skeleton_indices['smpl_24'] elif skeleton == b'coco_19': indices = self.per_skeleton_indices['coco_19'] elif skeleton == b'h36m_17': indices = self.per_skeleton_indices['h36m_17'] elif skeleton == b'h36m_25': indices = self.per_skeleton_indices['h36m_25'] elif skeleton == b'mpi_inf_3dhp_17': indices = self.per_skeleton_indices['mpi_inf_3dhp_17'] elif skeleton == b'mpi_inf_3dhp_28': indices = self.per_skeleton_indices['mpi_inf_3dhp_28'] elif skeleton == b'smpl+head_30': indices = self.per_skeleton_indices['smpl+head_30'] else: indices = tf.range(122, dtype=tf.int32) return tf.gather(poses, indices, axis=-2) @tf.function(input_signature=[ tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8), # image tf.TensorSpec(shape=(3, 3), dtype=tf.float32), # intrinsic_matrix tf.TensorSpec(shape=(5,), dtype=tf.float32), # distortion_coeffs tf.TensorSpec(shape=(4, 4), dtype=tf.float32), # extrinsic_matrix tf.TensorSpec(shape=(3,), dtype=tf.float32), # world_up_vector tf.TensorSpec(shape=(), dtype=tf.float32), # default_fov_degrees tf.TensorSpec(shape=(), dtype=tf.int32), # internal_batch_size tf.TensorSpec(shape=(), dtype=tf.int32), # antialias_factor tf.TensorSpec(shape=(), dtype=tf.int32), # num_aug tf.TensorSpec(shape=(), dtype=tf.bool), # average_aug tf.TensorSpec(shape=(), dtype=tf.string), # skeleton tf.TensorSpec(shape=(), dtype=tf.float32), # detector_threshold tf.TensorSpec(shape=(), dtype=tf.float32), # detector_nms_iou_threshold tf.TensorSpec(shape=(), dtype=tf.int32), # max_detections tf.TensorSpec(shape=(), dtype=tf.bool), # detector_flip_aug tf.TensorSpec(shape=(), dtype=tf.bool) # suppress_implausible_poses ]) def detect_poses( self, image, intrinsic_matrix=UNKNOWN_INTRINSIC_MATRIX, distortion_coeffs=DEFAULT_DISTORTION, extrinsic_matrix=DEFAULT_EXTRINSIC_MATRIX, world_up_vector=DEFAULT_WORLD_UP, default_fov_degrees=55, internal_batch_size=64, antialias_factor=1, num_aug=5, average_aug=True, skeleton='', detector_threshold=0.3, detector_nms_iou_threshold=0.7, max_detections=-1, detector_flip_aug=False, suppress_implausible_poses=True): images = image[tf.newaxis] intrinsic_matrix = intrinsic_matrix[tf.newaxis] distortion_coeffs = distortion_coeffs[tf.newaxis] extrinsic_matrix = extrinsic_matrix[tf.newaxis] result = self.detect_poses_batched( images, intrinsic_matrix, distortion_coeffs, extrinsic_matrix, world_up_vector, default_fov_degrees, internal_batch_size, antialias_factor, num_aug, average_aug, skeleton, detector_threshold, detector_nms_iou_threshold, max_detections, detector_flip_aug, suppress_implausible_poses) return tf.nest.map_structure(lambda x: tf.squeeze(x, 0), result) @tf.function(input_signature=[ tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8), # image tf.TensorSpec(shape=(None, 4), dtype=tf.float32), # boxes tf.TensorSpec(shape=(3, 3), dtype=tf.float32), # intrinsic_matrix tf.TensorSpec(shape=(5,), dtype=tf.float32), # distortion_coeffs tf.TensorSpec(shape=(4, 4), dtype=tf.float32), # extrinsic_matrix tf.TensorSpec(shape=(3,), dtype=tf.float32), # world_up_vector tf.TensorSpec(shape=(), dtype=tf.float32), # default_fov_degrees tf.TensorSpec(shape=(), dtype=tf.int32), # internal_batch_size tf.TensorSpec(shape=(), dtype=tf.int32), # antialias_factor tf.TensorSpec(shape=(), dtype=tf.int32), # num_aug tf.TensorSpec(shape=(), dtype=tf.bool), # average_aug tf.TensorSpec(shape=(), dtype=tf.string), # skeleton ]) def estimate_poses( self, image, boxes, intrinsic_matrix=UNKNOWN_INTRINSIC_MATRIX, distortion_coeffs=DEFAULT_DISTORTION, extrinsic_matrix=DEFAULT_EXTRINSIC_MATRIX, world_up_vector=DEFAULT_WORLD_UP, default_fov_degrees=55, internal_batch_size=64, antialias_factor=1, num_aug=5, average_aug=True, skeleton=''): images = image[tf.newaxis] boxes = tf.RaggedTensor.from_tensor(boxes[tf.newaxis]) intrinsic_matrix = intrinsic_matrix[tf.newaxis] distortion_coeffs = distortion_coeffs[tf.newaxis] extrinsic_matrix = extrinsic_matrix[tf.newaxis] result = self.estimate_poses_batched( images, boxes, intrinsic_matrix, distortion_coeffs, extrinsic_matrix, world_up_vector, default_fov_degrees, internal_batch_size, antialias_factor, num_aug, average_aug, skeleton) return tf.nest.map_structure(lambda x: tf.squeeze(x, 0), result)
return encoding_layer ################################################### # TF ENCODING HELPER FUNCTIONS ################################################### @tf.function def raged_lists_batch_to_multihot(ragged_lists_batch: tf.RaggedTensor, multihot_dim: int) -> tf.Tensor: """ Maps a batch of label indices to a batch of multi-hot ones """ # TODO: Seems tf.one_hot supports ragged tensors, so try to remove to_tensor call t = ragged_lists_batch.to_tensor(-1) # Default value = -1 -> one_hot will not assign any one t = tf.one_hot( t , multihot_dim ) t = tf.reduce_max( t , axis=1 ) return t @tf.function(input_signature=[tf.RaggedTensorSpec(shape=[None,None], dtype=tf.int64), tf.TensorSpec(shape=(), dtype=tf.bool)]) def pad_sequence_right( sequences_batch: tf.RaggedTensor, mask: bool) -> tf.Tensor: """ Pad sequences with zeros on right side """ # Avoid sequences larger than sequence_length: Get last sequence_length of each sequence sequences_batch = sequences_batch[:,-settings.settings.sequence_length:] if mask: # Add one to indices, to reserve 0 index for padding sequences_batch += 1 # Convert to dense, padding zeros to the right sequences_batch = sequences_batch.to_tensor(0, shape=[None, settings.settings.sequence_length]) return sequences_batch
def input_fn(self, data_file, n_repeat=1): batch_size = self.batch_size texts = list() mstr_sep_texts = list() all_label_strs, all_labels = list(), list() with open(data_file, encoding='utf-8') as f: for i, line in enumerate(f): x = json.loads(line) mstr = x['mention_span'] text = '{} [MASK] such as {} {}'.format( ' '.join(x['left_context_token']), mstr, ' '.join(x['right_context_token'])) # text = '{} {} {}'.format( # ' '.join(x['left_context_token']), mstr, ' '.join(x['right_context_token'])) mstr_sep_text = '{} [SEP] {} {} {}'.format( mstr, ' '.join(x['left_context_token']), mstr, ' '.join(x['right_context_token'])) # print(text) texts.append(text) mstr_sep_texts.append(mstr_sep_text) labels = x['y_str'] tids = [self.type_id_dict.get(t, -1) for t in labels] tids = [tid for tid in tids if tid > -1] # if i > 5: all_label_strs.append(labels) all_labels.append(tids) # if len(texts) >= 12: # break print(len(texts), 'texts') def tok_id_seq_gen(): text_ids, tok_id_seqs, tok_id_seqs_repeat = list(), list(), list() label_strs, y_vecs = list(), list() for _ in range(n_repeat): batch_id = 0 for text_idx, text in enumerate(texts): tokens = self.tokenizer.tokenize(text) tokens_full = ['[CLS]'] + tokens + ['[SEP]'] tok_id_seq = self.tokenizer.convert_tokens_to_ids( tokens_full) tok_id_seqs.append(tok_id_seq) fet_tokens = ['[CLS]'] + self.tokenizer.tokenize( mstr_sep_texts[text_idx]) + ['[SEP]'] fet_tok_id_seq = self.tokenizer.convert_tokens_to_ids( fet_tokens) # tok_id_seq = np.array([len(text)], np.float32) for _ in range(self.retriever_beam_size): tok_id_seqs_repeat.append(fet_tok_id_seq) text_ids.append(text_idx) y_vecs.append( to_one_hot(all_labels[text_idx], self.n_types)) label_strs.append(all_label_strs[text_idx]) if len(tok_id_seqs) >= batch_size: # tok_id_seq_batch, input_mask = get_padded_bert_input(tok_id_seqs) tok_id_seq_batch = tf.ragged.constant(tok_id_seqs) tok_id_seqs_repeat_ragged = tf.ragged.constant( tok_id_seqs_repeat) # y_vecs_tensor = tf.concat(y_vecs) # yield {'tok_id_seq_batch': tok_id_seq_batch, 'input_mask': input_mask}, y_vecs yield { 'batch_id': batch_id, 'tok_id_seq_batch': tok_id_seq_batch, 'tok_id_seqs_repeat': tok_id_seqs_repeat_ragged, 'text_ids': text_ids, 'labels': tf.constant(label_strs), 'tmp': tf.constant([[1, 2], [3, 4]], tf.int32), }, y_vecs text_ids, tok_id_seqs, tok_id_seqs_repeat, y_vecs = list( ), list(), list(), list() label_strs = list() batch_id += 1 # y_vecs = list() if len(tok_id_seqs) > 0: # tok_id_seq_batch, input_mask = get_padded_bert_input(tok_id_seqs) tok_id_seq_batch = tf.ragged.constant(tok_id_seqs) tok_id_seqs_repeat_ragged = tf.ragged.constant( tok_id_seqs_repeat) # y_vecs_tensor = tf.concat(y_vecs) # yield {'tok_id_seq_batch': tok_id_seq_batch, 'input_mask': input_mask}, y_vecs yield { 'batch_id': batch_id, 'tok_id_seq_batch': tok_id_seq_batch, 'tok_id_seqs_repeat': tok_id_seqs_repeat_ragged, 'text_ids': text_ids, 'labels': tf.constant(label_strs), 'tmp': tf.constant([[1, 2], [3, 4]], tf.int32), }, y_vecs # for v in iter(tok_id_seq_gen()): # print(v) dataset = tf.data.Dataset.from_generator( tok_id_seq_gen, output_signature=( { 'batch_id': tf.TensorSpec(shape=None, dtype=tf.int32), 'tok_id_seq_batch': tf.RaggedTensorSpec(dtype=tf.int32, ragged_rank=1), 'tok_id_seqs_repeat': tf.RaggedTensorSpec(dtype=tf.int32, ragged_rank=1), 'text_ids': tf.TensorSpec(shape=None, dtype=tf.int32), 'labels': tf.TensorSpec(shape=None, dtype=tf.string), 'tmp': tf.TensorSpec(shape=None, dtype=tf.int32), # 'block_emb': tf.TensorSpec(shape=block_emb_shape, dtype=tf.float32)}, }, tf.TensorSpec(shape=None, dtype=tf.float32))) return dataset
from tensorflow_transform import test_case _TEST_BATCH_SIZES = [1, 10] _TEST_DTYPES = [ tf.int16, tf.int32, tf.int64, tf.float32, tf.float64, tf.string, ] _TEST_TENSORS_TYPES = [ (lambda dtype: tf.TensorSpec([None], dtype=dtype), tf.Tensor, []), (lambda dtype: tf.TensorSpec([None, 2], dtype=dtype), tf.Tensor, [2]), (lambda dtype: tf.RaggedTensorSpec([None, None], dtype=dtype), tf.RaggedTensor, [None]), ( lambda dtype: tf.RaggedTensorSpec( # pylint: disable=g-long-lambda [None, None, 2], dtype=dtype, ragged_rank=1), tf.RaggedTensor, [None, 2]), ] class TF2UtilsTest(test_case.TransformTestCase): def test_strip_and_get_tensors_and_control_dependencies(self): @tf.function(input_signature=[tf.TensorSpec([], dtype=tf.int64)]) def func(x):
_RAGGED_TENSOR_TEST_CASES = (_MakeRaggedTensorDTypesTestCases() + [ dict(testcase_name="Simple", tensor_representation_textpb=""" ragged_tensor { feature_path { step: "ragged_feature" } row_partition_dtype: INT32 } """, record_batch=pa.RecordBatch.from_arrays([ pa.array([[1], None, [2], [3, 4, 5], []], type=pa.list_(pa.int64())) ], ["ragged_feature"]), expected_type_spec=tf.RaggedTensorSpec(tf.TensorShape([None, None]), tf.int64, ragged_rank=1, row_splits_dtype=tf.int32), expected_ragged_tensor=tf.compat.v1.ragged.RaggedTensorValue( values=np.asarray([1, 2, 3, 4, 5]), row_splits=np.asarray([0, 1, 1, 2, 5, 5]))), dict(testcase_name="3D", tensor_representation_textpb=""" ragged_tensor { feature_path { step: "ragged_feature" } row_partition_dtype: INT32 } """, record_batch=pa.RecordBatch.from_arrays([ pa.array([[[1]], None, [[2]], [[3, 4], [5]], []],
def type_to_tf_structure(type_spec: computation_types.Type): """Returns nested `tf.data.experimental.Structure` for a given TFF type. Args: type_spec: A `computation_types.Type`, the type specification must be composed of only named tuples and tensors. In all named tuples that appear in the type spec, all the elements must be named. Returns: An instance of `tf.data.experimental.Structure`, possibly nested, that corresponds to `type_spec`. Raises: ValueError: if the `type_spec` is composed of something other than named tuples and tensors, or if any of the elements in named tuples are unnamed. """ py_typecheck.check_type(type_spec, computation_types.Type) if type_spec.is_tensor(): return tf.TensorSpec(type_spec.shape, type_spec.dtype) elif type_spec.is_struct(): elements = structure.to_elements(type_spec) if not elements: return () element_outputs = [(k, type_to_tf_structure(v)) for k, v in elements] named = element_outputs[0][0] is not None if not all((e[0] is not None) == named for e in element_outputs): raise ValueError('Tuple elements inconsistently named.') if type_spec.python_container is None: if named: return collections.OrderedDict(element_outputs) else: return tuple(v for _, v in element_outputs) else: container_type = type_spec.python_container if (py_typecheck.is_named_tuple(container_type) or py_typecheck.is_attrs(container_type)): return container_type(**dict(element_outputs)) elif container_type is tf.RaggedTensor: flat_values = type_spec.flat_values nested_row_splits = type_spec.nested_row_splits ragged_rank = len(nested_row_splits) return tf.RaggedTensorSpec( shape=tf.TensorShape([None] * (ragged_rank + 1)), dtype=flat_values.dtype, ragged_rank=ragged_rank, row_splits_dtype=nested_row_splits[0].dtype, flat_values_spec=None) elif container_type is tf.SparseTensor: # We can't generally infer the shape from the type of the tensors, but # we *can* infer the rank based on the shapes of `indices` or # `dense_shape`. if (type_spec.indices.shape is not None and type_spec.indices.shape.dims[1] is not None): rank = type_spec.indices.shape.dims[1] shape = tf.TensorShape([None] * rank) elif (type_spec.dense_shape.shape is not None and type_spec.dense_shape.shape.dims[0] is not None): rank = type_spec.dense_shape.shape.dims[0] shape = tf.TensorShape([None] * rank) else: shape = None return tf.SparseTensorSpec(shape=shape, dtype=type_spec.values.dtype) elif named: return container_type(element_outputs) else: return container_type(e if e[0] is not None else e[1] for e in element_outputs) else: raise ValueError('Unsupported type {}.'.format( py_typecheck.type_string(type(type_spec))))
@author: landon """ import tensorflow as tf import numpy as np import json import argparse import os import time HAS_TITLE = [1,1,0,1,1,1,1,1,1] N_TRACKS = [0,5,5,10,25,25,100,100,1] IS_RANDOM = [0,0,0,0,0,1,0,1,0] TENSOR_SPEC = tf.RaggedTensorSpec(tf.TensorShape([4, None]), tf.int32, 1, tf.int64) def delete_tensor_by_indices(tensor,indices,n_tracks): idxs = tf.reshape(indices,(-1,1)) mask = ~tf.scatter_nd(indices=idxs,updates=tf.ones_like(indices,dtype=tf.bool),shape=[n_tracks]) return tf.boolean_mask(tensor,mask) @tf.autograph.experimental.do_not_convert def map_func(x): return { 'track_ids':x[0], 'title_ids':x[1], 'n_tracks':x[2][0], }
class UtilsTest(tf.test.TestCase, parameterized.TestCase): @parameterized.parameters( ((None, 3), tf.int32), ((2, ), tf.float64), ((2, None), tf.float32), ) def test_tensor_placeholder(self, shape, dtype): spec = tf.TensorSpec(shape=shape, dtype=dtype) placeholder = utils.placeholder(spec) assert tuple(placeholder.shape) == shape assert placeholder.dtype == dtype @parameterized.parameters( ((2, None), tf.int32, 1), ((None, None), tf.int64, 1), ((None, None, 2), tf.float64, 1), ((None, None, None), tf.float64, 2), ) def test_ragged_placeholder(self, shape, dtype, ragged_rank): expected = tf.RaggedTensorSpec(shape, ragged_rank=ragged_rank, dtype=dtype) placeholder = utils.placeholder(expected) assert_specs_equal(utils.type_spec(placeholder), expected) assert tuple(placeholder.shape) == shape assert placeholder.dtype == dtype @parameterized.parameters( ((None, None), tf.int32), ((None, 3), tf.int32), ((2, ), tf.float64), ((2, None), tf.float32), ) def test_sparse_placeholder(self, shape, dtype): spec = tf.SparseTensorSpec(shape, dtype) placeholder = utils.placeholder(spec) assert tuple(placeholder.shape) == shape assert placeholder.dtype == dtype # @parameterized.parameters( # (tf.TensorSpec((2,), tf.float64), 3, None, tf.TensorSpec, (3, 2)), # (tf.TensorSpec((None,), tf.float64), 3, False, tf.TensorSpec, (3, None)), # (tf.TensorSpec((None,), tf.float64), None, False, tf.TensorSpec, (None, None)), # ( # tf.TensorSpec((None,), tf.float64), # None, # True, # tf.RaggedTensorSpec, # (None, None), # ), # ( # tf.RaggedTensorSpec((None, None), tf.float32), # None, # None, # tf.RaggedTensorSpec, # (None, None, None), # ), # ( # tf.SparseTensorSpec((2, 3), tf.float64), # None, # None, # tf.SparseTensorSpec, # (None, 2, 3), # ), # ) # def test_batched_spec(self, spec, batch_size, ragged, expected_cls, expected_shape): # actual = utils.batched_spec(spec, batch_size=batch_size, ragged=ragged) # assert tuple(actual._shape) == expected_shape # assert isinstance(actual, expected_cls) # assert actual._dtype == spec._dtype @parameterized.parameters( ( tf.keras.Input((3, ), batch_size=2, dtype=tf.float64), tf.TensorSpec((2, 3), tf.float64), ), ( tf.keras.Input( shape=(None, ), batch_size=3, ragged=True, dtype=tf.float64), tf.RaggedTensorSpec((3, None), tf.float64), ), ( tf.keras.Input( shape=(4, ), batch_size=3, sparse=True, dtype=tf.float64), tf.SparseTensorSpec((3, 4), tf.float64), ), ) def test_type_spec(self, x, expected): assert_specs_equal(utils.type_spec(x), expected)
def call(self, ground_truth_boxes, anchors): # ground_truth_box [n_gt_boxes, box_dim] or [batch_size, n_gt_boxes, box_dim] # anchor [n_anchors, box_dim] def iou(ground_truth_box, anchor): # [n_anchors, 1] y_min_anchors, x_min_anchors, y_max_anchors, x_max_anchors = tf.split( anchor, num_or_size_splits=4, axis=-1) # [n_gt_boxes, 1] or [batch_size, n_gt_boxes, 1] y_min_gt, x_min_gt, y_max_gt, x_max_gt = tf.split( ground_truth_box, num_or_size_splits=4, axis=-1) # [n_anchors] anchor_areas = tf.squeeze((y_max_anchors - y_min_anchors) * (x_max_anchors - x_min_anchors), [1]) # [n_gt_boxes, 1] or [batch_size, n_gt_boxes, 1] gt_areas = (y_max_gt - y_min_gt) * (x_max_gt - x_min_gt) # [n_gt_boxes, n_anchors] or [batch_size, n_gt_boxes, n_anchors] max_y_min = tf.maximum(y_min_gt, tf.transpose(y_min_anchors)) min_y_max = tf.minimum(y_max_gt, tf.transpose(y_max_anchors)) intersect_heights = tf.maximum( tf.constant(0, dtype=ground_truth_box.dtype), (min_y_max - max_y_min)) # [n_gt_boxes, n_anchors] or [batch_size, n_gt_boxes, n_anchors] max_x_min = tf.maximum(x_min_gt, tf.transpose(x_min_anchors)) min_x_max = tf.minimum(x_max_gt, tf.transpose(x_max_anchors)) intersect_widths = tf.maximum( tf.constant(0, dtype=ground_truth_box.dtype), (min_x_max - max_x_min)) # [n_gt_boxes, n_anchors] or [batch_size, n_gt_boxes, n_anchors] intersections = intersect_heights * intersect_widths # [n_gt_boxes, n_anchors] or [batch_size, n_gt_boxes, n_anchors] unions = gt_areas + anchor_areas - intersections return tf.cast(tf.truediv(intersections, unions), tf.float32) if isinstance(ground_truth_boxes, tf.RaggedTensor): if anchors.shape.ndims == 2: return tf.map_fn( lambda x: iou(x, anchors), elems=ground_truth_boxes, parallel_iterations=32, back_prop=False, fn_output_signature=tf.RaggedTensorSpec(dtype=tf.float32, ragged_rank=0), ) else: return tf.map_fn( lambda x: iou(x[0], x[1]), elems=[ground_truth_boxes, anchors], parallel_iterations=32, back_prop=False, fn_output_signature=tf.RaggedTensorSpec(dtype=tf.float32, ragged_rank=0), ) if anchors.shape.ndims == 2: return iou(ground_truth_boxes, anchors) elif anchors.shape.ndims == 3: return tf.map_fn( lambda x: iou(x[0], x[1]), elems=[ground_truth_boxes, anchors], dtype=tf.float32, parallel_iterations=32, back_prop=False, )
class TensorToArrowTest(tf.test.TestCase, parameterized.TestCase): def _assert_tensor_alike_equal(self, left, right): self.assertIsInstance(left, type(right)) if isinstance(left, (tf.SparseTensor, tf.compat.v1.SparseTensorValue)): self.assertAllEqual(left.values, right.values) self.assertAllEqual(left.indices, right.indices) self.assertAllEqual(left.dense_shape, right.dense_shape) else: self.assertAllEqual(left, right) @parameterized.named_parameters(*_CONVERT_TEST_CASES) def test_convert( self, type_specs, expected_schema, expected_tensor_representations, tensor_input, expected_record_batch, options=tensor_to_arrow.TensorsToRecordBatchConverter.Options()): def convert_and_check(tensors, test_values_conversion): converter = tensor_to_arrow.TensorsToRecordBatchConverter( type_specs, options) self.assertEqual( {f.name: f.type for f in converter.arrow_schema()}, expected_schema, "actual: {}".format(converter.arrow_schema())) canonical_expected_tensor_representations = {} for n, r in expected_tensor_representations.items(): if not isinstance(r, schema_pb2.TensorRepresentation): r = text_format.Parse(r, schema_pb2.TensorRepresentation()) canonical_expected_tensor_representations[n] = r self.assertEqual(canonical_expected_tensor_representations, converter.tensor_representations()) rb = converter.convert(tensors) self.assertLen(expected_record_batch, rb.num_columns) for i, column in enumerate(rb): expected = expected_record_batch[rb.schema[i].name] self.assertTrue( column.equals(expected), "{}: actual: {}, expected: {}".format( rb.schema[i].name, column, expected)) # Test that TensorAdapter(TensorsToRecordBatchConverter()) is identity. adapter = tensor_adapter.TensorAdapter( tensor_adapter.TensorAdapterConfig( arrow_schema=converter.arrow_schema(), tensor_representations=converter.tensor_representations())) adapter_output = adapter.ToBatchTensors( rb, produce_eager_tensors=not test_values_conversion) self.assertEqual(adapter_output.keys(), tensors.keys()) for k in adapter_output.keys(): if "value" not in k: self._assert_tensor_alike_equal(adapter_output[k], tensors[k]) def ragged_tensor_to_value(tensor): if isinstance(tensor, tf.RaggedTensor): values = tensor.values return tf.compat.v1.ragged.RaggedTensorValue( values=ragged_tensor_to_value(values), row_splits=tensor.row_splits.numpy()) else: return tensor.numpy() def convert_eager_to_value(tensor): if isinstance(tensor, tf.SparseTensor): return tf.compat.v1.SparseTensorValue(tensor.indices, tensor.values, tensor.dense_shape) elif isinstance(tensor, tf.Tensor): return tensor.numpy() elif isinstance(tensor, tf.RaggedTensor): return ragged_tensor_to_value(tensor) else: raise NotImplementedError( "Only support converting SparseTensors, Tensors and RaggedTensors. " "Got: {}".format(type(tensor))) if tf.__version__ >= "2": convert_and_check(tensor_input, test_values_conversion=False) if tf.executing_eagerly(): values_input = { k: convert_eager_to_value(v) for k, v in tensor_input.items() } else: some_tensor = next(iter(tensor_input.values())) graph = (some_tensor.row_splits[0].graph if isinstance( some_tensor, tf.RaggedTensor) else some_tensor.graph) with tf.compat.v1.Session(graph=graph) as s: values_input = s.run(tensor_input) convert_and_check(values_input, test_values_conversion=True) def test_relaxed_varlen_sparse_tensor(self): # Demonstrates that TensorAdapter(TensorsToRecordBatchConverter()) is not # an identity if the second dense dimension of SparseTensor is not tight. type_specs = {"sp": tf.SparseTensorSpec([None, None], tf.int32)} sp = tf.compat.v1.SparseTensorValue(values=np.array([1, 2], np.int32), indices=[[0, 0], [2, 0]], dense_shape=[4, 2]) if tf.__version__ >= "2": sp = tf.SparseTensor.from_value(sp) converter = tensor_to_arrow.TensorsToRecordBatchConverter(type_specs) rb = converter.convert({"sp": sp}) adapter = tensor_adapter.TensorAdapter( tensor_adapter.TensorAdapterConfig( arrow_schema=converter.arrow_schema(), tensor_representations=converter.tensor_representations())) adapter_output = adapter.ToBatchTensors( rb, produce_eager_tensors=tf.__version__ >= "2") self.assertAllEqual(sp.values, adapter_output["sp"].values) self.assertAllEqual(sp.indices, adapter_output["sp"].indices) self.assertAllEqual(adapter_output["sp"].dense_shape, [4, 1]) def test_unable_to_handle(self): with self.assertRaisesRegex(ValueError, "No handler found"): tensor_to_arrow.TensorsToRecordBatchConverter( {"sp": tf.SparseTensorSpec([None, None, None], tf.int32)}) with self.assertRaisesRegex(ValueError, "No handler found"): tensor_to_arrow.TensorsToRecordBatchConverter( {"sp": tf.SparseTensorSpec([None, None], tf.bool)}) def test_incompatible_type_spec(self): converter = tensor_to_arrow.TensorsToRecordBatchConverter( {"sp": tf.SparseTensorSpec([None, None], tf.int32)}) sp_cls = (tf.SparseTensor if tf.__version__ >= "2" else tf.compat.v1.SparseTensorValue) with self.assertRaisesRegex(TypeError, "Expected SparseTensorSpec"): converter.convert({ "sp": sp_cls(indices=[[0, 1]], values=tf.constant([0], dtype=tf.int64), dense_shape=[4, 1]) }) @parameterized.named_parameters(*[ dict(testcase_name="bool_value_type", spec=tf.RaggedTensorSpec(shape=[2, None, None], dtype=tf.bool, ragged_rank=2, row_splits_dtype=tf.int64)), dict(testcase_name="2d_leaf_value", spec=tf.RaggedTensorSpec(shape=[2, None, None], dtype=tf.int32, ragged_rank=1, row_splits_dtype=tf.int64)), dict(testcase_name="ragged_rank_less_than_one", spec=tf.RaggedTensorSpec(shape=[2], dtype=tf.int32, ragged_rank=0, row_splits_dtype=tf.int64)), ]) def test_unable_to_handle_ragged(self, spec): with self.assertRaisesRegex(ValueError, "No handler found"): tensor_to_arrow.TensorsToRecordBatchConverter({"rt": spec})
def _build_tree_from_leaf(leaf_nodes: tf.Tensor, arity: int) -> tf.RaggedTensor: """A function constructs a complete tree given all the leaf nodes. The function takes a 1-D array representing the leaf nodes of a tree and the tree's arity, and constructs a complete tree by recursively summing the adjacent children to get the parent until reaching the root node. Because we assume a complete tree, if the number of leaf nodes does not divide arity, the leaf nodes will be padded with zeros. Args: leaf_nodes: A 1-D array storing the leaf nodes of the tree. arity: A `int` for the branching factor of the tree, i.e. the number of children for each internal node. Returns: `tf.RaggedTensor` representing the tree. For example, if `leaf_nodes=tf.Tensor([1, 2, 3, 4])` and `arity=2`, then the returned value should be `tree=tf.RaggedTensor([[10],[3,7],[1,2,3,4]])`. In this way, `tree[layer][index]` can be used to access the node indexed by (layer, index) in the tree, """ def pad_zero(leaf_nodes, size): paddings = tf.zeros( shape=(size - leaf_nodes.shape[0],), dtype=leaf_nodes.dtype) return tf.concat((leaf_nodes, paddings), axis=0) leaf_nodes_size = tf.constant(leaf_nodes.shape[0], dtype=tf.float32) num_layers = tf.math.ceil( tf.math.log(leaf_nodes_size) / tf.math.log(tf.cast(arity, dtype=tf.float32))) + 1 leaf_nodes = pad_zero( leaf_nodes, tf.math.pow(tf.cast(arity, dtype=tf.float32), num_layers - 1)) def _shrink_layer(layer: tf.Tensor, arity: int) -> tf.Tensor: return tf.reduce_sum((tf.reshape(layer, (-1, arity))), 1) # The following `tf.while_loop` constructs the tree from bottom up by # iteratively applying `_shrink_layer` to each layer of the tree. The reason # for the choice of TF1.0-style `tf.while_loop` is that @tf.function does not # support auto-translation from python loop to tf loop when loop variables # contain a `RaggedTensor` whose shape changes across iterations. idx = tf.identity(num_layers) loop_cond = lambda i, h: tf.less_equal(2.0, i) def _loop_body(i, h): return [ tf.add(i, -1.0), tf.concat(([_shrink_layer(h[0], arity)], h), axis=0) ] _, tree = tf.while_loop( loop_cond, _loop_body, [idx, tf.RaggedTensor.from_tensor([leaf_nodes])], shape_invariants=[ idx.get_shape(), tf.RaggedTensorSpec(dtype=leaf_nodes.dtype, ragged_rank=1) ]) return tree
# # all_vals = get_ragged_vals(5) # for length in lens: # vals = list() # for _ in range(length): # vals.append(random.uniform(-1, 1)) # blocks_list.append(vals) blocks_list = [[3, 4], [5, 3, 8], [8]] def data_gen(): for vals in blocks_list: yield tf.ragged.constant([vals], dtype=tf.float32) dataset = tf.data.Dataset.from_generator( data_gen, output_signature=tf.RaggedTensorSpec(ragged_rank=1, dtype=tf.float32)) for v in dataset: print(v) # dataset = dataset.apply( # tf.data.experimental.dense_to_ragged_batch(2)) # for v in dataset.batch(4): # print(v) # # print(v[1].to_tensor()) # vt = v[1].to_tensor() # print(vt) # print(tf.squeeze(vt)) # for v in dataset: # print(v) # tf.data.experimental.save(dataset, '/data/hldai/data/tmp/tmp.tfdata', shard_func=lambda x: np.int64(0))
def input_fn(): import json from locbert import tokenization batch_size = 4 data_file = '/data/hldai/data/ultrafine/uf_data/crowd/test.json' type_vocab_file = '/data/hldai/data/ultrafine/uf_data/ontology/types.txt' reader_module_path = '/data/hldai/data/realm_data/cc_news_pretrained/bert' vocab_file = os.path.join(reader_module_path, 'assets/vocab.txt') tokenizer = tokenization.FullTokenizer(vocab_file, do_lower_case=True) types, type_id_dict = datautils.load_vocab_file(type_vocab_file) n_types = len(types) # texts = ['He is a teacher.', # 'He teaches his students.', # 'He is a lawyer.'] texts = list() all_labels = list() with open(data_file, encoding='utf-8') as f: for i, line in enumerate(f): x = json.loads(line) text = '{} {} {}'.format( ' '.join(x['left_context_token']), x['mention_span'], ' '.join(x['right_context_token'])) # print(text) texts.append(text) labels = x['y_str'] tids = [type_id_dict.get(t, -1) for t in labels] tids = [tid for tid in tids if tid > -1] # if i > 5: all_labels.append(tids) if len(texts) >= 8: break print(len(texts), 'texts') def tok_id_seq_gen(): tok_id_seqs = list() y_vecs = list() for i, text in enumerate(texts): tokens = tokenizer.tokenize(text) # print(tokens) tokens_full = ['[CLS]'] + tokens + ['[SEP]'] tok_id_seq = tokenizer.convert_tokens_to_ids(tokens_full) # tok_id_seq = np.array([len(text)], np.float32) tok_id_seqs.append(tok_id_seq) y_vecs.append(to_one_hot(all_labels[i], n_types)) if len(tok_id_seqs) >= batch_size: # tok_id_seq_batch, input_mask = get_padded_bert_input(tok_id_seqs) tok_id_seq_batch = tf.ragged.constant(tok_id_seqs) # y_vecs_tensor = tf.concat(y_vecs) yield {'tok_id_seq_batch': tok_id_seq_batch, # 'input_mask': input_mask, 'vals': np.random.uniform(-1, 1, (3, 5))}, y_vecs tok_id_seqs = list() y_vecs = list() if len(tok_id_seqs) > 0: # tok_id_seq_batch, input_mask = get_padded_bert_input(tok_id_seqs) # y_vecs_tensor = tf.concat(y_vecs) tok_id_seq_batch = tf.ragged.constant(tok_id_seqs) yield {'tok_id_seq_batch': tok_id_seq_batch, # 'input_mask': input_mask, 'vals': np.random.uniform(-1, 1, (3, 5))}, y_vecs # for v in iter(tok_id_seq_gen()): # print(v) dataset = tf.data.Dataset.from_generator( tok_id_seq_gen, output_signature=( { 'tok_id_seq_batch': tf.RaggedTensorSpec(dtype=tf.int32, ragged_rank=1), # 'tok_id_seq_batch': tf.TensorSpec(shape=None, dtype=tf.int32), # 'input_mask': tf.TensorSpec(shape=None, dtype=tf.int32), 'vals': tf.TensorSpec(shape=None, dtype=tf.float32) }, tf.TensorSpec(shape=None, dtype=tf.float32))) return dataset
class Pose3dEstimator(tf.Module): def __init__(self, model_path, antialias_factor=4): super().__init__() self.antialias_factor = antialias_factor self.crop_model = tf.saved_model.load(model_path) self.crop_side = 256 self.joint_names = self.crop_model.joint_names self.joint_edges = self.crop_model.joint_edges joint_names = [b.decode('utf8') for b in self.joint_names.numpy()] self.joint_info = data.datasets3d.JointInfo(joint_names, self.joint_edges.numpy()) self.__call__.get_concrete_function( tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8), tf.TensorSpec(shape=(None, 3, 3), dtype=tf.float32), tf.RaggedTensorSpec(shape=(None, None, 4), ragged_rank=1, dtype=tf.float32), tf.TensorSpec(shape=(), dtype=tf.int32), tf.TensorSpec(shape=(), dtype=tf.int32)) self.__call__.get_concrete_function( tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8), tf.TensorSpec(shape=(3, 3), dtype=tf.float32), tf.TensorSpec(shape=(None, 4), dtype=tf.float32), tf.TensorSpec(shape=(), dtype=tf.int32), tf.TensorSpec(shape=(), dtype=tf.int32)) @tf.function def __call__(self, image, intrinsic_matrix, boxes, internal_batch_size=64, n_aug=5): if image.shape.rank == 3: return self.predict_single_image( image, intrinsic_matrix, boxes, internal_batch_size, n_aug) else: return self.predict_multi_image( image, intrinsic_matrix, boxes, internal_batch_size, n_aug) @tf.function(input_signature=[ tf.TensorSpec(shape=(None, None, 3), dtype=tf.uint8), tf.TensorSpec(shape=(3, 3), dtype=tf.float32), tf.TensorSpec(shape=(None, 4), dtype=tf.float32), tf.TensorSpec(shape=(), dtype=tf.int32), tf.TensorSpec(shape=(), dtype=tf.int32)]) def predict_single_image(self, image, intrinsic_matrix, boxes, internal_batch_size=64, n_aug=5): if tf.size(boxes) == 0: return tf.zeros(shape=(0, self.joint_info.n_joints, 3)) ragged_boxes = tf.RaggedTensor.from_tensor(boxes[np.newaxis]) return self.predict_multi_image( image[np.newaxis], intrinsic_matrix[np.newaxis], ragged_boxes, internal_batch_size, n_aug)[0] @tf.function(input_signature=[ tf.TensorSpec(shape=(None, None, None, 3), dtype=tf.uint8), tf.TensorSpec(shape=(None, 3, 3), dtype=tf.float32), tf.RaggedTensorSpec(shape=(None, None, 4), ragged_rank=1, dtype=tf.float32), tf.TensorSpec(shape=(), dtype=tf.int32), tf.TensorSpec(shape=(), dtype=tf.int32)]) def predict_multi_image(self, image, intrinsic_matrix, boxes, internal_batch_size=64, n_aug=5): """Estimate 3D human poses in camera space for multiple bounding boxes specified for an image. """ n_images = tf.shape(image)[0] if tf.size(boxes) == 0: # Special case for zero boxes provided result_flat = tf.zeros(shape=(0, self.joint_info.n_joints, 3)) return tf.RaggedTensor.from_row_lengths(result_flat, tf.zeros(n_images, tf.int64)) boxes_flat = boxes.flat_values n_box_per_image = boxes.row_lengths() image_id_per_box = boxes.value_rowids() n_total_boxes = tf.shape(boxes_flat)[0] # image = (tf.cast(image, tf.float32) / np.float32(255))# ** 2.2 if tf.shape(intrinsic_matrix)[0] == 1: intrinsic_matrix = tf.repeat(intrinsic_matrix, n_images, axis=0) Ks = tf.repeat(intrinsic_matrix, n_box_per_image, axis=0) gammas = tf.cast(tf.linspace(0.6, 1.0, n_aug), tf.float16) angle_range = np.float32(np.deg2rad(25)) angles = tf.linspace(-angle_range, angle_range, n_aug) scales = tf.concat([ linspace_noend(0.8, 1.0, n_aug // 2), tf.linspace(1.0, 1.1, n_aug - n_aug // 2)], axis=0) flips = (tf.range(n_aug) - n_aug // 2) % 2 != 0 flipmat = tf.constant([[-1, 0, 0], [0, 1, 0], [0, 0, 1]], np.float32) maybe_flip = tf.where(flips[:, np.newaxis, np.newaxis], flipmat, tf.eye(3)) rotmat = rotation_mat_zaxis(-angles) crops_per_batch = internal_batch_size // n_aug if crops_per_batch == 0: # No batching results_flat = self.predict_single_batch( image, Ks, boxes_flat, image_id_per_box, n_aug, rotmat, maybe_flip, flips, scales, gammas) return tf.RaggedTensor.from_row_lengths(results_flat, n_box_per_image) n_batches = tf.cast(tf.math.ceil(n_total_boxes / crops_per_batch), tf.int32) result_batches = tf.TensorArray( tf.float32, size=n_batches, element_shape=(None, self.joint_info.n_joints, 3), infer_shape=False) for i in tf.range(n_batches): box_batch = boxes_flat[i * crops_per_batch:(i + 1) * crops_per_batch] image_ids = image_id_per_box[i * crops_per_batch:(i + 1) * crops_per_batch] K_batch = Ks[i * crops_per_batch:(i + 1) * crops_per_batch] poses = self.predict_single_batch( image, K_batch, box_batch, image_ids, n_aug, rotmat, maybe_flip, flips, scales, gammas) result_batches = result_batches.write(i, poses) results_flat = result_batches.concat() return tf.RaggedTensor.from_row_lengths(results_flat, n_box_per_image) def predict_single_batch( self, images, K, boxes, image_ids, n_aug, rotmat, maybe_flip, flips, scales, gammas): n_box = tf.shape(boxes)[0] center_points = boxes[:, :2] + boxes[:, 2:4] / 2 box_center_camspace = transf(center_points - K[:, :2, 2], tf.linalg.inv(K[:, :2, :2])) box_center_camspace = tf.concat( [box_center_camspace, tf.ones_like(box_center_camspace[:, :1])], axis=1) new_z = box_center_camspace / tf.linalg.norm(box_center_camspace, axis=-1, keepdims=True) new_x = tf.stack([new_z[:, 2], tf.zeros_like(new_z[:, 2]), -new_z[:, 0]], axis=1) new_y = tf.linalg.cross(new_z, new_x) nonaug_R = tf.stack([new_x, new_y, new_z], axis=1) new_R = maybe_flip[:, np.newaxis] @ rotmat[:, np.newaxis] @ nonaug_R box_scales = self.crop_side / tf.reduce_max(boxes[:, 2:4], axis=-1) new_K_mid = (tf.reshape(scales, [-1, 1, 1, 1]) * tf.reshape(box_scales, [1, -1, 1, 1]) * tf.reshape(K[:, :2, :2], [1, -1, 2, 2])) intrinsic_matrix = tf.concat([ tf.concat([new_K_mid, tf.fill((n_aug, n_box, 2, 1), self.crop_side / 2)], axis=3), tf.concat([tf.zeros((n_aug, n_box, 1, 2), tf.float32), tf.ones((n_aug, n_box, 1, 1), tf.float32)], axis=3)], axis=2) new_proj_matrix = intrinsic_matrix @ new_R homography = K @ tf.linalg.inv(new_proj_matrix) intrinsic_matrix_flat = tf.reshape(intrinsic_matrix, [n_aug * n_box, 3, 3]) homography = tf.reshape(homography, [n_aug, n_box, 3, 3]) homography = tf.reshape(homography, [-1, 9]) homography = homography[:, :8] / homography[:, 8:] if self.antialias_factor > 1: H = homography a = self.antialias_factor homography = tf.stack([H[:, 0] / a, H[:, 1] / a, H[:, 2] - (a - 1) / 2, H[:, 3] / a, H[:, 4] / a, H[:, 5] - (a - 1) / 2, H[:, 6] / a, H[:, 7] / a], axis=1) temp_side = self.crop_side * self.antialias_factor image_ids = tf.tile(image_ids, [n_aug]) crops = perspective_transform( images, homography, (temp_side, temp_side), 'BILINEAR', image_ids) crops = tf.cast(crops, tf.float16) / 255 if self.antialias_factor > 1: crops = tf.image.resize( crops, (self.crop_side, self.crop_side), method=tf.image.ResizeMethod.AREA, antialias=True) crops = tf.reshape(crops, [n_aug, n_box * self.crop_side, self.crop_side, 3]) crops **= tf.reshape(gammas, [-1, 1, 1, 1]) crops = tf.reshape(crops, [-1, self.crop_side, self.crop_side, 3]) poses = self.crop_model(crops, intrinsic_matrix_flat) poses = tf.reshape(poses, [n_aug, -1, tf.shape(poses)[1], 3]) left_right_swapped = tf.gather(poses, self.joint_info.mirror_mapping, axis=2) poses = tf.where(tf.reshape(flips, [-1, 1, 1, 1]), left_right_swapped, poses) poses_origspace = tf.einsum('...nk,...jk->...nj', poses, tf.linalg.inv(new_R)) return tf.reduce_mean(poses_origspace, axis=0)
class Prediction(tf.Module): NOT_FOUND_INDEX = -1 def __init__(self, model=None): self.item_labels_path = Labels.item_labels_path() self.customer_labels_path = Labels.customer_labels_path() if model: self.model: tf.keras.Model = model else: self.model: tf.keras.Model = tf.keras.models.load_model( settings.get_model_path('exported_model')) #print(">>>", self.model) self.model.summary() # Lookup table: Customer label -> Customer index self.customer_labels_lookup = tf.lookup.StaticHashTable( tf.lookup.TextFileInitializer(self.customer_labels_path, tf.string, tf.lookup.TextFileIndex.WHOLE_LINE, tf.int64, tf.lookup.TextFileIndex.LINE_NUMBER, delimiter=" "), Prediction.NOT_FOUND_INDEX) # Lookup table: Item label -> Item index self.item_labels_lookup = tf.lookup.StaticHashTable( tf.lookup.TextFileInitializer(self.item_labels_path, tf.string, tf.lookup.TextFileIndex.WHOLE_LINE, tf.int64, tf.lookup.TextFileIndex.LINE_NUMBER, delimiter=" "), Prediction.NOT_FOUND_INDEX) # Reverse lookup tables (item index -> item string label) self.item_indices_lookup = tf.lookup.StaticHashTable( tf.lookup.TextFileInitializer(self.item_labels_path, tf.int64, tf.lookup.TextFileIndex.LINE_NUMBER, tf.string, tf.lookup.TextFileIndex.WHOLE_LINE, delimiter=" "), "") @tf.function( input_signature=[tf.TensorSpec(shape=[None, None], dtype=tf.int32)]) def post_process_items(self, batch_items_indices: tf.Tensor) -> tf.Tensor: # Do lookup return self.item_indices_lookup.lookup( tf.cast(batch_items_indices, tf.int64)) @tf.function(input_signature=[ tf.TensorSpec(shape=[None, None], dtype=tf.float32), tf.TensorSpec(shape=[], dtype=tf.int64) ]) def _top_predictions_tensor(self, results, n_results): # Get most probable item indices sorted_indices = tf.argsort(results, direction='DESCENDING') top_indices = sorted_indices[:, 0:n_results] top_probabilities = tf.gather(results, top_indices, batch_dims=1) # Convert item indices to item labels #print("top_indices", top_indices) top_item_labels = self.post_process_items(top_indices) return top_item_labels, top_probabilities @staticmethod @tf.function(input_signature=[ tf.RaggedTensorSpec(shape=[None, None], dtype=tf.int64), tf.TensorSpec(shape=[None, None], dtype=tf.float32) ]) def _remove_input_items_from_prediction(batch_item_indices, result): # batch_item_indices is a ragged with input indices for each batch row. Ex [ [0] , [1, 2] ] batch_item_indices = batch_item_indices.to_tensor( -1) # -> [ [0,-1] , [1, 2] ] #print(batch_item_indices, batch_item_indices.shape[0]) # Convert batch_item_indices row indices to (row,column) indices row_indices = tf.range(0, tf.shape(batch_item_indices)[0], dtype=tf.int64) # -> [ 0, 1 ] row_indices = tf.repeat( row_indices, [tf.shape(batch_item_indices)[1]]) # -> [ 0, 0, 1, 1 ] #print(">>>", batch_item_indices) batch_item_indices = tf.reshape(batch_item_indices, shape=[-1]) # -> [ 0, -1, 1, 2 ] batch_item_indices = tf.stack( [row_indices, batch_item_indices], axis=1) # -> [ [0,0] , [0,-1], [1,1], [1,2] ] # batch_item_indices.to_tensor(-1) added -1's to pad the matrix. Remove these indices # Needed according to tf.tensor_scatter_nd_update doc. (it will fail in CPU execution, if there are out of bound indices) # Get indices without -1's: gather_idxs = tf.where( batch_item_indices[:, 1] != -1) # -> [[0], [2], [3]] batch_item_indices = tf.gather_nd( batch_item_indices, gather_idxs) # -> [ [0,0] , [1,1], [1,2] ] # To remove input indices, we will set a probability -1 in their indices updates = tf.repeat( -1.0, tf.shape(batch_item_indices)[0]) # -> [ -1, -1, -1 ] # Assign -1's to the input indices: return tf.tensor_scatter_nd_update(result, batch_item_indices, updates) @staticmethod @tf.function(input_signature=[ tf.RaggedTensorSpec(shape=[None, None], dtype=tf.int64), tf.TensorSpec(shape=(), dtype=tf.int64) ]) def count_non_equal(batch: tf.RaggedTensor, value: tf.Tensor) -> tf.Tensor: """ Returns the count of elements in 'batch'' distincts to 'value' on each batch row """ elements_equal_to_value = tf.not_equal(batch, value) as_ints = tf.cast(elements_equal_to_value, tf.int64) count = tf.reduce_sum(as_ints, axis=1) return count @staticmethod @tf.function(input_signature=[ tf.RaggedTensorSpec(shape=[None, None], dtype=tf.int64), tf.TensorSpec(shape=(), dtype=tf.int64) ]) def remove_not_found_index(batch_item_indices: tf.RaggedTensor, not_found_index: tf.Tensor) -> tf.RaggedTensor: # Count non -1's on each row found_counts_per_row = Prediction.count_non_equal( batch_item_indices, not_found_index) #print("found_counts_per_row", found_counts_per_row) # Get non -1 values batch_item_indices from flat values flat_values = batch_item_indices.flat_values mask = tf.not_equal(flat_values, not_found_index) #print("mask", mask ) flat_found_indices = tf.boolean_mask(flat_values, mask) #print("flat_found_indices", flat_found_indices ) return tf.RaggedTensor.from_row_lengths(flat_found_indices, found_counts_per_row) @tf.function(input_signature=[ tf.RaggedTensorSpec(shape=[None, None], dtype=tf.string) ]) def preprocess_items(self, batch_item_labels: tf.RaggedTensor): #print("batch_item_labels ->", batch_item_labels) # Do lookups item label -> index, -1 if not found batch_item_indices = tf.ragged.map_flat_values( self.item_labels_lookup.lookup, batch_item_labels) #print( "batch_item_indices", batch_item_indices ) # Remove -1's: batch_item_indices = Prediction.remove_not_found_index( batch_item_indices, Prediction.NOT_FOUND_INDEX) #print( "batch_item_indices >>>", batch_item_indices ) # Remove duplicated items # TODO: UNIMPLEMENTED. tf.unique works only with 1D dimensions... # batch_item_indices = tf.map_fn(lambda x: tf.unique(x), batch_item_indices.to_tensor(-1) ) # print("unique", tf.unique(batch_item_indices)) return batch_item_indices @tf.function def prepreprocess_customers(self, batch_customer_labels: tf.Tensor): # Get customer label batch_customer_indices = self.customer_labels_lookup.lookup( batch_customer_labels) # Get the "UNKNOWN" customer index unknown_customer_index = self.customer_labels_lookup.lookup( tf.constant(Labels.UNKNOWN_LABEL, dtype=tf.string)) # Replace -1 by "UNKNOWN" index update_indices = tf.where( tf.math.equal(batch_customer_indices, Prediction.NOT_FOUND_INDEX)) batch_customer_indices = tf.tensor_scatter_nd_update( batch_customer_indices, update_indices, tf.repeat(unknown_customer_index, tf.size(update_indices))) return batch_customer_indices @tf.function def _run_model_and_postprocess(self, batch_item_indices, batch_customer_indices, n_results): # Run the model batch = (batch_item_indices, batch_customer_indices) result = self.model(batch, training=False) # Convert logits to probabilities result = tf.nn.softmax(result) if settings.model_type == ModelType.GPT: # GPT return probabilities for each sequence timestep. # We need the probabilities for the LAST input timested. # Batch is a ragged tensors (ex. [[1], [2,3]]), so, get the probabilities for last element position in each # batch sequence: indices = batch_item_indices.row_lengths() indices -= 1 #print("indices", indices) result = tf.gather(result, indices, batch_dims=1) #print("probs", probs) # Set result[ batch_item_indices ] = -1.0: #print("batch_item_indices >>>***", batch_item_indices) result = Prediction._remove_input_items_from_prediction( batch_item_indices, result) # Get most probable n results return self._top_predictions_tensor(result, n_results) @tf.function def _run_model_filter_empty_sequences(self, batch_item_indices: tf.RaggedTensor, batch_customer_indices, n_results): # Check if there are empty sequences sequences_lenghts = batch_item_indices.row_lengths() non_empty_seq_count = tf.math.count_nonzero(sequences_lenghts) n_sequences = tf.shape(sequences_lenghts, tf.int64)[0] #print(">>>", non_empty_seq_count, n_results) if non_empty_seq_count == 0: # All sequences are empty label_predictions = tf.zeros([n_sequences, n_results], dtype=tf.string) probs_predictions = tf.zeros([n_sequences, n_results], dtype=tf.float32) return (label_predictions, probs_predictions) elif non_empty_seq_count >= n_sequences: # There are no empty sequences. Run the model with the full batch return self._run_model_and_postprocess(batch_item_indices, batch_customer_indices, n_results) else: # There are some empty sequences # Model will fail if a sequence is empty, and it seems it's the expected behaviour: Do not feed empty sequences # Get non empty sequences mask non_empty_mask = tf.math.greater(sequences_lenghts, 0) # Get non empty sequences non_empty_sequences: tf.RaggedTensor = tf.ragged.boolean_mask( batch_item_indices, non_empty_mask) non_empty_customers = tf.boolean_mask(batch_customer_indices, non_empty_mask) # Run model label_predictions, probs_predictions = self._run_model_and_postprocess( non_empty_sequences, non_empty_customers, n_results) # Merge real predictions with empty predictions for empty sequences: indices = tf.where(non_empty_mask) final_shape = [n_sequences, n_results] label_predictions = tf.scatter_nd(indices, label_predictions, final_shape) #print(label_predictions) probs_predictions = tf.scatter_nd(indices, probs_predictions, final_shape) #print(probs_predictions) return (label_predictions, probs_predictions) @tf.function(input_signature=[ tf.RaggedTensorSpec(shape=[None, None], dtype=tf.string), tf.TensorSpec(shape=[None], dtype=tf.string), tf.TensorSpec(shape=[], dtype=tf.int64) ]) def run_model_prediction(self, batch_item_labels, batch_customer_labels, n_results): # Convert labels to indices #print(">>> batch_item_labels", batch_item_labels) batch_item_indices = self.preprocess_items(batch_item_labels) #print(">>> batch_item_indices", batch_item_indices) #print(">>> batch_customer_labels", batch_customer_labels) batch_customer_indices = self.prepreprocess_customers( batch_customer_labels) #print(">>> batch_customer_indices", batch_customer_indices) # Run the model return self._run_model_filter_empty_sequences(batch_item_indices, batch_customer_indices, n_results) @tf.function(input_signature=[ tf.TensorSpec(shape=[None], dtype=tf.string), tf.TensorSpec(shape=[], dtype=tf.string), tf.TensorSpec(shape=[], dtype=tf.int64) ]) def run_model_single(self, item_labels, customer_label, n_results): # Convert single example to batch batch_item_labels = tf.expand_dims(item_labels, axis=0) ragged_batch_item_labels = tf.RaggedTensor.from_tensor( batch_item_labels) batch_customer_labels = tf.expand_dims(customer_label, axis=0) predicted_item_labels, predicted_item_probs = self.run_model_prediction( ragged_batch_item_labels, batch_customer_labels, n_results) # Remove batch dimension predicted_item_labels = tf.squeeze(predicted_item_labels, 0) predicted_item_probs = tf.squeeze(predicted_item_probs, 0) return (predicted_item_labels, predicted_item_probs) def predict_batch(self, transactions: List[Transaction], n_items_result: int) -> List: # Setup batch batch = Transaction.to_net_inputs_batch(transactions) #print("*** batch", batch) batch = (tf.ragged.constant(batch[0], dtype=tf.string), tf.constant(batch[1])) #print("*** batch", batch) results = self.run_model_prediction(batch[0], batch[1], n_items_result) #print("raw", results) results = (results[0].numpy(), results[1].numpy()) return results def predict_single(self, transaction: Transaction, n_items_result: int) -> List[Tuple[str, float]]: results = self.predict_batch([transaction], n_items_result) return (results[0][0], results[1][0])
class Prediction(tf.Module): """ Run and process candidates generation model predictions """ # Directory in "models/[MODEL]/" dir where to export the model CHECKPOINTS_DIR = 'checkpoints' # Directory in "models/[MODEL]/" dir where to export the model EXPORTED_MODEL_DIR = 'exported_model' def __init__(self, model: tf.keras.Model = None): if model: self.model: tf.keras.Model = model else: self.model: tf.keras.Model = tf.keras.models.load_model( settings.settings.get_model_path( False, Prediction.EXPORTED_MODEL_DIR)) #self.model.summary() @tf.function(input_signature=[ tf.TensorSpec(shape=[None, None], dtype=tf.float32), tf.TensorSpec(shape=[], dtype=tf.int64) ]) def _top_predictions_tensor(self, results, n_results) -> Tuple[tf.Tensor, tf.Tensor]: """ Returns a tuple (items indices, items probabilities) with most probable items, up to "n_results" """ # Get most probable item indices sorted_indices = tf.argsort(results, direction='DESCENDING') top_indices = sorted_indices[:, 0:n_results] top_probabilities = tf.gather(results, top_indices, batch_dims=1) return top_indices, top_probabilities @staticmethod @tf.function(input_signature=[ tf.RaggedTensorSpec(shape=[None, None], dtype=tf.int64), tf.TensorSpec(shape=[None, None], dtype=tf.float32) ]) def _remove_input_items_from_prediction(batch_item_indices, result): """ Remove input items from predictions, as we don't want to predict them. It is done setting their probabilities to -1 Args: batch_item_indices: Batch input item indices result: Batch predicted probabilities Returns: Batch predicted probabilities with input items probs. set to -1 """ # batch_item_indices is a ragged with input indices for each batch row. Ex [ [0] , [1, 2] ] batch_item_indices = batch_item_indices.to_tensor( -1) # -> [ [0,-1] , [1, 2] ] #print(batch_item_indices, batch_item_indices.shape[0]) # Convert batch_item_indices row indices to (row,column) indices row_indices = tf.range(0, tf.shape(batch_item_indices)[0], dtype=tf.int64) # -> [ 0, 1 ] row_indices = tf.repeat( row_indices, [tf.shape(batch_item_indices)[1]]) # -> [ 0, 0, 1, 1 ] #print(">>>", batch_item_indices) batch_item_indices = tf.reshape(batch_item_indices, shape=[-1]) # -> [ 0, -1, 1, 2 ] batch_item_indices = tf.stack( [row_indices, batch_item_indices], axis=1) # -> [ [0,0] , [0,-1], [1,1], [1,2] ] # batch_item_indices.to_tensor(-1) added -1's to pad the matrix. Remove these indices # Needed according to tf.tensor_scatter_nd_update doc. (it will fail in CPU execution, if there are out of bound indices) # Get indices without -1's: gather_idxs = tf.where( batch_item_indices[:, 1] != -1) # -> [[0], [2], [3]] batch_item_indices = tf.gather_nd( batch_item_indices, gather_idxs) # -> [ [0,0] , [1,1], [1,2] ] # To remove input indices, we will set a probability -1 in their indices updates = tf.repeat( -1.0, tf.shape(batch_item_indices)[0]) # -> [ -1, -1, -1 ] # Assign -1's to the input indices: return tf.tensor_scatter_nd_update(result, batch_item_indices, updates) @tf.function def _run_model_and_postprocess(self, inputs_batch, item_indices_feature_name, n_results): # Run the model result = self.model(inputs_batch, training=False) # Convert logits to probabilities result = tf.nn.softmax(result) # Label indices feature values batch_item_indices = inputs_batch[item_indices_feature_name] if settings.settings.model_type == settings.ModelType.GPT: # GPT return probabilities for each sequence timestep. # We need the probabilities for the LAST input timested. # Batch is a ragged tensors (ex. [[1], [2,3]]), so, get the probabilities for last element position in each # batch sequence: indices = batch_item_indices.row_lengths() indices -= 1 #print("indices", indices) result = tf.gather(result, indices, batch_dims=1) #print("probs", probs) # Set result[ batch_item_indices ] = -1.0: result = Prediction._remove_input_items_from_prediction( batch_item_indices, result) # Get most probable n results return self._top_predictions_tensor(result, n_results) def _preprocess_transactions( self, transactions: List[Transaction] ) -> Tuple[List[Transaction], List[int]]: """ Remove unknown items and empty transactions. Returns: Non empty transactions, with labels replaced by indices, and indices of empty transactions in the original transactions list """ result = [], empty_sequences_idxs = [] for idx, trn in enumerate(transactions): trn = trn.replace_labels_by_indices() # Now, trn item sequence contains -1 for unknown items. Remove sequence feature for these items trn = trn.remove_unknown_item_indices() if trn.sequence_length() == 0: # Empty sequence. If we feed it to the model, it can fail (rnn layers) empty_sequences_idxs.append(idx) else: result.append(trn) return result, empty_sequences_idxs def _transactions_to_model_inputs( self, transactions: List[Transaction]) -> Dict[str, tf.Tensor]: # Prepare dictionary with features names inputs_dict = { feature.name: [] for feature in settings.settings.features } # Concatenate transaction values for each feature for trn in transactions: # Use labels indices instead raw values for feature in settings.settings.features: inputs_dict[feature.name].append(trn[feature.name]) # To tensor values for feature in settings.settings.features: if feature.sequence: inputs_dict[feature.name] = tf.ragged.constant( inputs_dict[feature.name], dtype=tf.int64) else: inputs_dict[feature.name] = tf.constant( inputs_dict[feature.name], dtype=tf.int64) # Keras inputs are mapped by position, so return result as a list #return [ inputs_dict[feature.name] for feature in settings.settings.features ] return inputs_dict def predict_raw_batch( self, transactions: List[Transaction], n_items_result: int ) -> Tuple[np.ndarray, np.ndarray, List[tf.Tensor]]: # Transactions to keras inputs batch = self._transactions_to_model_inputs(transactions) # Run prediction if len(batch) > 0: top_item_indices, top_probabilities = self._run_model_and_postprocess( batch, settings.settings.features.item_label_feature, n_items_result) top_item_indices = top_item_indices.numpy() top_probabilities = top_probabilities.numpy() else: top_item_indices = np.array([], dtype=int) top_probabilities = np.array([], dtype=float) return top_item_indices, top_probabilities, batch def predict_batch(self, transactions: List[Transaction], n_items_result: int, preprocess=True) -> Tuple[np.ndarray, np.ndarray]: if preprocess: # Convert labels to indices. Remove unknown items and empty sequences transactions, empty_sequences_idxs = self._preprocess_transactions( transactions) else: empty_sequences_idxs = [] top_item_indices, top_probabilities = self.predict_raw_batch( transactions, n_items_result) # Convert item indices to labels top_item_labels = settings.settings.features.items_sequence_feature( ).labels.indices_to_labels(top_item_indices) # Insert fake results for empty sequences if len(empty_sequences_idxs) > 0: empty_labels = np.zeros([n_items_result], dtype=str) empty_probs = np.zeros([n_items_result], dtype=float) # Inserts cannot be done all at same time, np.insert expects indices to the array before insert, and we don't have them for idx in empty_sequences_idxs: top_item_labels = np.insert(top_item_labels, idx, empty_labels, axis=0) top_probabilities = np.insert(top_probabilities, idx, empty_probs, axis=0) return top_item_labels, top_probabilities def predict_single(self, transaction: Transaction, n_items_result: int, preprocess=True): top_item_labels, top_probabilities = self.predict_batch([transaction], n_items_result, preprocess) return top_item_labels[0], top_probabilities[0]
class EAMPotential(tf.keras.Model): """Base class for all EAM based potentials""" def __init__(self, atom_types, build_forces=False, preprocessed_input=False, cutoff=10., method='partition_stitch', force_method='old'): """ atom_types: list of str used to construct all the model layers. Note the order is important as it determines the corresponding integer representation that is fed to the model. build_forces: boolean determines if the graph for evaluating forces should be constructed. preprocessed_input: boolean switches between input of atomic positions and interatomic distances. Preprocessed is currently faster as the calculation of interatomic distances in the model can not be parallelized. cutoff: float if preprocessed_input=False this determines the global cutoff until which interatomic distances are calculated. method: str switch for different implementations of the main body options are: - 'partition_stitch' default and probably savest choice - 'where' uses branchless programming can be significantly faster as it is easier to parallelize - 'gather_scatter' old implementation, no longer maintained force_method: str switch for debugging, options are: - 'old' default - 'new' """ super().__init__() self.atom_types = atom_types self.type_dict = {} for i, t in enumerate(atom_types): self.type_dict[t] = i self.atom_pair_types = [] self.pair_type_dict = {} for i, (t1, t2) in enumerate( combinations_with_replacement(self.atom_types, 2)): t = ''.join([t1, t2]) self.atom_pair_types.append(t) self.pair_type_dict[t] = i self.build_forces = build_forces self.preprocessed_input = preprocessed_input self.method = method self.force_method = force_method inputs = { 'types': tf.keras.Input(shape=(None, 1), ragged=True, dtype=tf.int32) } if self.preprocessed_input: inputs['pair_types'] = tf.keras.Input(shape=(None, None, 1), ragged=True, dtype=inputs['types'].dtype) inputs['distances'] = tf.keras.Input(shape=(None, None, 1), ragged=True) if self.build_forces: # [batchsize] x N x (N-1) x N x 3 inputs['dr_dx'] = tf.keras.Input(shape=(None, None, None, 3), ragged=True) else: self.cutoff = cutoff inputs['positions'] = tf.keras.Input(shape=(None, 3), ragged=True) self._set_inputs(inputs) (self.pair_potentials, self.pair_rho, self.embedding_functions, self.offsets) = self.build_functions() def build_functions(self): pass def call(self, inputs): if self.preprocessed_input: return self.shallow_call(inputs) return self.deep_call(inputs) def deep_call(self, inputs): types = inputs['types'] positions = inputs['positions'] distances, pair_types, dr_dx = tf.map_fn( lambda x: distances_and_pair_types(x[0], x[1], len( self.atom_types), self.cutoff), (positions, types), fn_output_signature=(tf.RaggedTensorSpec(shape=[None, None, 1], ragged_rank=1, dtype=positions.dtype), tf.RaggedTensorSpec(shape=[None, None, 1], ragged_rank=1, dtype=types.dtype), tf.RaggedTensorSpec( shape=[None, None, None, 3], ragged_rank=2, dtype=positions.dtype))) if self.build_forces: return self.main_body_with_forces(types, distances, pair_types, dr_dx) return self.main_body_no_forces(types, distances, pair_types) @tf.function def shallow_call(self, inputs): types = inputs['types'] distances = inputs['distances'] pair_types = inputs['pair_types'] if self.build_forces: dr_dx = inputs['dr_dx'] return self.main_body_with_forces(types, distances, pair_types, dr_dx) return self.main_body_no_forces(types, distances, pair_types) @tf.function(input_signature=(tf.RaggedTensorSpec( tf.TensorShape([None, None, 1]), tf.int32, 1, tf.int64), tf.RaggedTensorSpec( tf.TensorShape([None, None, None, 1]), tf.keras.backend.floatx(), 2, tf.int64), tf.RaggedTensorSpec( tf.TensorShape([None, None, None, 1]), tf.int32, 2, tf.int64))) def main_body_no_forces(self, types, distances, pair_types): """Calculates the energy per atom by calling the main body """ if self.method == 'partition_stitch': energy = self.body_partition_stitch(types, distances, pair_types) elif self.method == 'gather_scatter': energy = self.body_gather_scatter(types, distances, pair_types) elif self.method == 'where': energy = self.body_where(types, distances, pair_types) else: raise NotImplementedError('Unknown method %s' % self.method) number_of_atoms = tf.cast(types.row_lengths(), energy.dtype, name='number_of_atoms') energy_per_atom = tf.divide(energy, tf.expand_dims(number_of_atoms, axis=-1), name='energy_per_atom') return {'energy_per_atom': energy_per_atom} @tf.function( input_signature=(tf.RaggedTensorSpec(tf.TensorShape([None, None, 1]), tf.int32, 1, tf.int64), tf.RaggedTensorSpec( tf.TensorShape([None, None, None, 1]), tf.keras.backend.floatx(), 2, tf.int64), tf.RaggedTensorSpec( tf.TensorShape([None, None, None, 1]), tf.int32, 2, tf.int64), tf.RaggedTensorSpec( tf.TensorShape([None, None, None, None, 3]), tf.keras.backend.floatx(), 3, tf.int64))) def main_body_with_forces(self, types, distances, pair_types, dr_dx): """Calculates the energy per atom and the derivative of the total energy with respect to the distances """ with tf.GradientTape() as tape: tape.watch(distances.flat_values) if self.method == 'partition_stitch': energy = self.body_partition_stitch(types, distances, pair_types) elif self.method == 'gather_scatter': energy = self.body_gather_scatter(types, distances, pair_types) elif self.method == 'where': energy = self.body_where(types, distances, pair_types) else: raise NotImplementedError('Unknown method %s' % self.method) number_of_atoms = tf.cast(types.row_lengths(), energy.dtype, name='number_of_atoms') energy_per_atom = tf.divide(energy, tf.expand_dims(number_of_atoms, axis=-1), name='energy_per_atom') if self.force_method == 'old': # Probably not should not be reshaped to RaggedTensor only to sum # over these dimensions in the next step. dE_dr = tf.RaggedTensor.from_nested_row_splits( tape.gradient(energy, distances.flat_values), distances.nested_row_splits, name='dE_dr') # dr_dx.shape = (batch_size, None, None, None, 3) # dE_dr.shape = (batch_size, None, None, 1) # Sum over atom indices i and j. Force is the negative gradient. forces = -tf.reduce_sum(dr_dx * tf.expand_dims(dE_dr, -1), axis=(-3, -4), name='dE_dr_times_dr_dx') elif self.force_method == 'new': dr_dx = dr_dx.merge_dims(1, 2) print(dr_dx.nested_row_splits) dE_dr = tf.RaggedTensor.from_row_splits(tape.gradient( energy, distances.flat_values), dr_dx.row_splits, name='dE_dr') # dr_dx.shape = (batch_size, None, None, 3) # dE_dr.shape = (batch_size, None, 1) forces = -tf.reduce_sum(dr_dx * tf.expand_dims(dE_dr, -1), axis=-3, name='dE_dr_times_dr_dx') else: raise NotImplementedError('Unknown force method %s' % self.force_method) return {'energy_per_atom': energy_per_atom, 'forces': forces} @tf.function def body_where(self, types, distances, pair_types): """""" rho = distances phi = distances for t in self.atom_pair_types: cond = tf.equal(pair_types, self.pair_type_dict[t]) rho = ragged_where( cond, tf.ragged.map_flat_values(self.pair_rho[t], rho), rho) phi = ragged_where( cond, tf.ragged.map_flat_values(self.pair_potentials[t], phi), phi) # Sum over atoms j sum_rho = tf.reduce_sum(rho**2, axis=-2, name='sum_rho') sum_phi = tf.reduce_sum(phi, axis=-2, name='sum_phi') # Make sure that sum_rho is never exactly zero since this leads to # problems in the gradient of the square root embedding function embedding_energies = tf.math.maximum(sum_rho, 1e-30) for t in self.atom_types: cond = tf.equal(types, self.type_dict[t]) # tf.abs(x) necessary here since most embedding functions do not # support negative inputs but embedding energies are typically # negative and tf.where applies the function to every vector entry embedding_energies = ragged_where( cond, tf.ragged.map_flat_values( lambda x: (self.embedding_functions[t] (tf.abs(x)) + self.offsets[t](x)), embedding_energies), embedding_energies) atomic_energies = sum_phi + embedding_energies # Sum over atoms i return tf.reduce_sum(atomic_energies, axis=-2, name='energy') @tf.function def body_partition_stitch(self, types, distances, pair_types): """main body using dynamic_partition and dynamic_stitch methods""" pair_type_indices = tf.dynamic_partition(tf.expand_dims( tf.range(tf.size(distances)), -1), pair_types.flat_values, len(self.atom_pair_types), name='pair_type_indices') # Partition distances according to the pair_type partitioned_r = tf.dynamic_partition(distances.flat_values, pair_types.flat_values, len(self.atom_pair_types), name='partitioned_r') rho = [ self.pair_rho[t](tf.expand_dims(part, -1)) for t, part in zip(self.atom_pair_types, partitioned_r) ] phi = [ self.pair_potentials[t](tf.expand_dims(part, -1)) for t, part in zip(self.atom_pair_types, partitioned_r) ] rho = tf.dynamic_stitch(pair_type_indices, rho) phi = tf.dynamic_stitch(pair_type_indices, phi) # Reshape to ragged tensors phi = tf.RaggedTensor.from_nested_row_splits( phi, distances.nested_row_splits) rho = tf.RaggedTensor.from_nested_row_splits( rho, distances.nested_row_splits) # Sum over atoms j sum_rho = tf.reduce_sum(rho**2, axis=-2, name='sum_rho') sum_phi = tf.reduce_sum(phi, axis=-2, name='sum_phi') # Make sure that sum_rho is never exactly zero since this leads to # problems in the gradient of the square root embedding function sum_rho = tf.math.maximum(sum_rho, 1e-30) # Embedding energy partitioned_sum_rho = tf.dynamic_partition(sum_rho, types, len(self.atom_types), name='partitioned_sum_rho') type_indices = tf.dynamic_partition(tf.expand_dims( tf.range(tf.size(sum_rho)), -1), types.flat_values, len(self.atom_types), name='type_indices') # energy offset is added here embedding_energies = [ self.embedding_functions[t](tf.expand_dims(rho_t, -1)) + self.offsets[t](tf.expand_dims(rho_t, -1)) for t, rho_t in zip(self.atom_types, partitioned_sum_rho) ] embedding_energies = tf.dynamic_stitch(type_indices, embedding_energies, name='embedding_energies') atomic_energies = sum_phi.flat_values + embedding_energies # Reshape to ragged atomic_energies = tf.RaggedTensor.from_row_splits( atomic_energies, types.row_splits, name='atomic_energies') # Sum over atoms i return tf.reduce_sum(atomic_energies, axis=-2, name='energy') @tf.function def body_gather_scatter(self, types, distances, pair_types): """main body using gather and scatter methods""" phi = tf.zeros_like(distances.flat_values) rho = tf.zeros_like(distances.flat_values) for ij, type_ij in enumerate(self.atom_pair_types): # Flat values necessary until where properly supports # ragged tensors indices = tf.where(tf.equal(pair_types, ij).flat_values)[:, 0] masked_distances = tf.gather(distances.flat_values, indices) phi = tf.tensor_scatter_nd_update( phi, tf.expand_dims(indices, -1), self.pair_potentials[type_ij](masked_distances)) rho = tf.tensor_scatter_nd_update( rho, tf.expand_dims(indices, -1), self.pair_rho[type_ij](masked_distances)) # Reshape back to ragged tensors phi = tf.RaggedTensor.from_nested_row_splits( phi, distances.nested_row_splits) rho = tf.RaggedTensor.from_nested_row_splits( rho, distances.nested_row_splits) # Sum over atoms j and flatten again atomic_energies = tf.reduce_sum(phi, axis=-2).flat_values sum_rho = tf.reduce_sum(rho**2, axis=-2).flat_values # Make sure that sum_rho is never exactly zero since this leads to # problems in the gradient of the square root embedding function sum_rho = tf.math.maximum(sum_rho, 1e-30) for i, t in enumerate(self.atom_types): indices = tf.where(tf.equal(types, i).flat_values)[:, 0] atomic_energies = tf.tensor_scatter_nd_add( atomic_energies, tf.expand_dims(indices, -1), self.embedding_functions[t](tf.gather(sum_rho, indices))) # Reshape to ragged atomic_energies = tf.RaggedTensor.from_row_splits( atomic_energies, types.row_splits) # Sum over atoms i return tf.reduce_sum(atomic_energies, axis=-2, name='energy') def tabulate(self, filename, atomic_numbers, atomic_masses, lattice_constants=dict(), lattice_types=dict(), cutoff_rho=120.0, nrho=10000, cutoff=6.2, nr=10000): """ Uses the atsim module to tabulate the potential for use in LAMMPS or similar programs. filename: str path of the output file. Note that the extension .eam.fs is added atomic_numbers: dict containing the atomic_numbers of all self.atom_types atomic_masses: dict containing the atomic mass of all self.atom_types """ from atsim.potentials import EAMPotential, Potential from atsim.potentials.eam_tabulation import SetFL_FS_EAMTabulation warnings.warn('Long range behavior of the tabulated potentials differs' ' from the tensorflow implementation!') def pair_wrapper(fun): def wrapped_fun(x): return 2 * fun(tf.reshape(x, (1, 1))) return wrapped_fun def rho_wrapper(fun): def wrapped_fun(x): return tf.math.square(fun(tf.reshape(x, (1, 1)))) return wrapped_fun def F_wrapper(fun): def wrapped_fun(x): return fun(tf.reshape(x, (1, 1))) return wrapped_fun pair_potentials = [] pair_densities = {t: {} for t in self.atom_types} for (t1, t2) in combinations_with_replacement(self.atom_types, 2): pair_type = ''.join([t1, t2]) def pair_pot(x): return self.pair_potentials[pair_type](tf.reshape(x, (1, 1))) pair_potentials.append( Potential( t1, t2, pair_wrapper(self.pair_potentials[pair_type]) #lambda x: self.pair_potentials[pair_type]( # x*tf.ones((1, 1))) )) pair_densities[t1][t2] = rho_wrapper(self.pair_rho[pair_type]) pair_densities[t2][t1] = rho_wrapper(self.pair_rho[pair_type]) eam_potentials = [] for t in self.atom_types: eam_potentials.append( EAMPotential(t, atomic_numbers[t], atomic_masses[t], F_wrapper(self.embedding_functions[t]), pair_densities[t], latticeConstant=lattice_constants.get(t, 0.0), latticeType=lattice_types.get(t, 'fcc'))) tabulation = SetFL_FS_EAMTabulation(pair_potentials, eam_potentials, cutoff, nr, cutoff_rho, nrho) with open(''.join([filename, '.eam.fs']), 'w') as outfile: tabulation.write(outfile)