import functools import operator import string import hypothesis.strategies as st import numpy as np import numpy.testing as npt from caffe2.python import core, dataset, workspace from caffe2.python.dataset import Const from caffe2.python.schema import ( FeedRecord, FetchRecord, Field, List, Map, NewRecord, Scalar, Struct, from_blob_list, ) from caffe2.python.test_util import TestCase from hypothesis import given def _assert_arrays_equal(actual, ref, err_msg): if ref.dtype.kind in ("S", "O", "U"): np.testing.assert_array_equal(actual, ref, err_msg=err_msg) else: np.testing.assert_allclose(actual, ref, atol=1e-4, rtol=1e-4, err_msg=err_msg) def _assert_records_equal(actual, ref): assert isinstance(actual, Field) assert isinstance(ref, Field) b1 = actual.field_blobs() b2 = ref.field_blobs() assert len(b1) == len(b2), "Records have different lengths: %d vs. %d" % ( len(b1), len(b2), ) for name, d1, d2 in zip(ref.field_names(), b1, b2): _assert_arrays_equal(d1, d2, err_msg="Mismatch in field %s." % name) @st.composite def _sparse_features_map(draw, num_records, **kwargs): sparse_maps_lengths = draw( st.lists( st.integers(min_value=1, max_value=10), min_size=num_records, max_size=num_records, ) ) sparse_maps_total_length = sum(sparse_maps_lengths) sparse_keys = draw( st.lists( st.integers(min_value=1, max_value=100), min_size=sparse_maps_total_length, max_size=sparse_maps_total_length, unique=True, ) ) sparse_values_lengths = draw( st.lists( st.integers(min_value=1, max_value=10), min_size=sparse_maps_total_length, max_size=sparse_maps_total_length, ) ) total_sparse_values_lengths = sum(sparse_values_lengths) sparse_values = draw( # max_value is max int64 st.lists( st.integers(min_value=1, max_value=9223372036854775807), min_size=total_sparse_values_lengths, max_size=total_sparse_values_lengths, ) ) return [ sparse_maps_lengths, sparse_keys, sparse_values_lengths, sparse_values, ] @st.composite def _dense_features_map(draw, num_records, **kwargs): float_lengths = draw( st.lists( st.integers(min_value=1, max_value=10), min_size=num_records, max_size=num_records, ) ) total_length = sum(float_lengths) float_keys = draw( st.lists( st.integers(min_value=1, max_value=100), min_size=total_length, max_size=total_length, unique=True, ) ) float_values = draw( st.lists(st.floats(), min_size=total_length, max_size=total_length) ) return [float_lengths, float_keys, float_values] @st.composite def _dataset(draw, min_elements=3, max_elements=10, **kwargs): schema = Struct( # Dense Features Map ("floats", Map(Scalar(np.int32), Scalar(np.float32))), # Sparse Features Map ( "int_lists", Map( Scalar(np.int32), List(Scalar(np.int64)), ), ), # Complex Type ("text", Scalar(str)), ) num_records = draw(st.integers(min_value=min_elements, max_value=max_elements)) raw_dense_features_map_contents = draw(_dense_features_map(num_records)) raw_sparse_features_map_contents = draw(_sparse_features_map(num_records)) raw_text_contents = [ draw( st.lists( st.text(alphabet=string.ascii_lowercase), min_size=num_records, max_size=num_records, ) ) ] # Concatenate all raw contents to a single one contents_raw = ( raw_dense_features_map_contents + raw_sparse_features_map_contents + raw_text_contents ) contents = from_blob_list(schema, contents_raw) return (schema, contents, num_records) class TestDatasetOps(TestCase): @given(_dataset()) def test_pack_unpack(self, input): """ Tests if packing and unpacking of the whole dataset is an identity. """ (schema, contents, num_records) = input dataset_fields = schema.field_names() for pack_to_single_shared_ptr in (True, False): net = core.Net("pack_unpack_net") batch = NewRecord(net, contents) FeedRecord(batch, contents) packed = net.PackRecords( batch.field_blobs(), 1, fields=dataset_fields, pack_to_single_shared_ptr=pack_to_single_shared_ptr, ) unpacked = packed.UnPackRecords( [], len(dataset_fields), fields=dataset_fields ) workspace.RunNetOnce(net) for initial_tensor, unpacked_tensor in zip(batch.field_blobs(), unpacked): npt.assert_array_equal( workspace.FetchBlob(initial_tensor), workspace.FetchBlob(unpacked_tensor), ) def test_dataset_ops(self): """ 1. Defining the schema of our dataset. This example schema could represent, for example, a search query log. """ schema = Struct( # fixed size vector, which will be stored as a matrix when batched ("dense", Scalar((np.float32, 3))), # could represent a feature map from feature ID to float value ("floats", Map(Scalar(np.int32), Scalar(np.float32))), # could represent a multi-valued categorical feature map ( "int_lists", Map( Scalar(np.int32), List(Scalar(np.int64)), ), ), # could represent a multi-valued, weighted categorical feature map ( "id_score_pairs", Map( Scalar(np.int32), Map( Scalar(np.int64), Scalar(np.float32), keys_name="ids", values_name="scores", ), ), ), # additional scalar information ( "metadata", Struct( ("user_id", Scalar(np.int64)), ("user_embed", Scalar((np.float32, 2))), ("query", Scalar(str)), ), ), ) """ This is what the flattened fields for this schema look like, along with its type. Each one of these fields will be stored, read and written as a tensor. """ expected_fields = [ ("dense", (np.float32, 3)), ("floats:lengths", np.int32), ("floats:values:keys", np.int32), ("floats:values:values", np.float32), ("int_lists:lengths", np.int32), ("int_lists:values:keys", np.int32), ("int_lists:values:values:lengths", np.int32), ("int_lists:values:values:values", np.int64), ("id_score_pairs:lengths", np.int32), ("id_score_pairs:values:keys", np.int32), ("id_score_pairs:values:values:lengths", np.int32), ("id_score_pairs:values:values:values:ids", np.int64), ("id_score_pairs:values:values:values:scores", np.float32), ("metadata:user_id", np.int64), ("metadata:user_embed", (np.float32, 2)), ("metadata:query", str), ] zipped = zip(expected_fields, schema.field_names(), schema.field_types()) for (ref_name, ref_type), name, dtype in zipped: self.assertEquals(ref_name, name) self.assertEquals(np.dtype(ref_type), dtype) """ 2. The contents of our dataset. Contents as defined below could represent, for example, a log of search queries along with dense, sparse features and metadata. The dataset below has 3 top-level entries. """ contents_raw = [ # dense [[1.1, 1.2, 1.3], [2.1, 2.2, 2.3], [3.1, 3.2, 3.3]], # floats [1, 2, 3], # len [11, 21, 22, 31, 32, 33], # key [1.1, 2.1, 2.2, 3.1, 3.2, 3.3], # value # int lists [2, 0, 1], # len [11, 12, 31], # key [2, 4, 3], # value:len [111, 112, 121, 122, 123, 124, 311, 312, 313], # value:value # id score pairs [1, 2, 2], # len [11, 21, 22, 31, 32], # key [1, 1, 2, 2, 3], # value:len [111, 211, 221, 222, 311, 312, 321, 322, 323], # value:ids [11.1, 21.1, 22.1, 22.2, 31.1, 31.2, 32.1, 32.2, 32.3], # val:score # metadata [123, 234, 456], # user_id [[0.2, 0.8], [0.5, 0.5], [0.7, 0.3]], # user_embed ["dog posts", "friends who like to", "posts about ca"], # query ] # convert the above content to ndarrays, checking against the schema contents = from_blob_list(schema, contents_raw) """ 3. Creating and appending to the dataset. We first create an empty dataset with the given schema. Then, a Writer is used to append these entries to the dataset. """ ds = dataset.Dataset(schema) net = core.Net("init") with core.NameScope("init"): ds.init_empty(net) content_blobs = NewRecord(net, contents) FeedRecord(content_blobs, contents) writer = ds.writer(init_net=net) writer.write_record(net, content_blobs) workspace.RunNetOnce(net) """ 4. Iterating through the dataset contents. If we were to iterate through the top level entries of our dataset, this is what we should expect to see: """ entries_raw = [ ( [[1.1, 1.2, 1.3]], # dense [1], [11], [1.1], # floats [2], [11, 12], [2, 4], [111, 112, 121, 122, 123, 124], # intlst [1], [11], [1], [111], [11.1], # id score pairs [123], [[0.2, 0.8]], ["dog posts"], # metadata ), ( [[2.1, 2.2, 2.3]], # dense [2], [21, 22], [2.1, 2.2], # floats [0], [], [], [], # int list [2], [21, 22], [1, 2], [211, 221, 222], [21.1, 22.1, 22.2], [234], [[0.5, 0.5]], ["friends who like to"], # metadata ), ( [[3.1, 3.2, 3.3]], # dense [3], [31, 32, 33], [3.1, 3.2, 3.3], # floats [1], [31], [3], [311, 312, 313], # int lst [2], [31, 32], [2, 3], [311, 312, 321, 322, 323], [31.1, 31.2, 32.1, 32.2, 32.3], # id score list [456], [[0.7, 0.3]], ["posts about ca"], # metadata ), # after the end of the dataset, we will keep getting empty vectors ([],) * 16, ([],) * 16, ] entries = [from_blob_list(schema, e) for e in entries_raw] """ Let's go ahead and create the reading nets. We will run `read` net multiple times and assert that we are reading the entries the way we stated above. """ read_init_net = core.Net("read_init") read_next_net = core.Net("read_next") reader = ds.reader(read_init_net) should_continue, batch = reader.read_record(read_next_net) workspace.RunNetOnce(read_init_net) workspace.CreateNet(read_next_net, True) for entry in entries: workspace.RunNet(str(read_next_net)) actual = FetchRecord(batch) _assert_records_equal(actual, entry) """ 5. Reading/writing in a single plan If all of operations on the data are expressible as Caffe2 operators, we don't need to load the data to python, iterating through the dataset in a single Plan. Where we will process the dataset a little and store it in a second dataset. We can reuse the same Reader since it supports reset. """ reset_net = core.Net("reset_net") reader.reset(reset_net) read_step, batch = reader.execution_step() """ We will add the line number * 1000 to the feature ids. """ process_net = core.Net("process") line_no = Const(process_net, 0, dtype=np.int32) const_one = Const(process_net, 1000, dtype=np.int32) process_net.Add([line_no, const_one], [line_no]) field = batch.floats.keys.get() process_net.Print(field, []) process_net.Add([field, line_no], field, broadcast=1, axis=0) """ Lets create a second dataset and append to it. """ ds2 = dataset.Dataset(schema, name="dataset2") ds2.init_empty(reset_net) writer = ds2.writer(reset_net) writer.write_record(process_net, batch) # commit is not necessary for DatasetWriter but will add it for # generality of the example commit_net = core.Net("commit") writer.commit(commit_net) """ Time to create and run a plan which will do the processing """ plan = core.Plan("process") plan.AddStep(core.execution_step("reset", reset_net)) plan.AddStep(read_step.AddNet(process_net)) plan.AddStep(core.execution_step("commit", commit_net)) workspace.RunPlan(plan) """ Now we should have dataset2 populated. """ ds2_data = FetchRecord(ds2.content()) field = ds2_data.floats.keys field.set(blob=field.get() - [1000, 2000, 2000, 3000, 3000, 3000]) _assert_records_equal(contents, ds2_data) """ 6. Slicing a dataset You can create a new schema from pieces of another schema and reuse the same data. """ subschema = Struct(("top_level", schema.int_lists.values)) int_list_contents = contents.int_lists.values.field_names() self.assertEquals(len(subschema.field_names()), len(int_list_contents)) """ 7. Random Access a dataset """ read_init_net = core.Net("read_init") read_next_net = core.Net("read_next") idx = np.array([2, 1, 0]) indices_blob = Const(read_init_net, idx, name="indices") reader = ds.random_reader(read_init_net, indices_blob) reader.computeoffset(read_init_net) should_stop, batch = reader.read_record(read_next_net) workspace.CreateNet(read_init_net, True) workspace.RunNetOnce(read_init_net) workspace.CreateNet(read_next_net, True) for i in range(len(entries)): k = idx[i] if i in idx else i entry = entries[k] workspace.RunNet(str(read_next_net)) actual = FetchRecord(batch) _assert_records_equal(actual, entry) workspace.RunNet(str(read_next_net)) self.assertEquals(True, workspace.FetchBlob(should_stop)) """ 8. Random Access a dataset with loop_over = true """ read_init_net = core.Net("read_init") read_next_net = core.Net("read_next") idx = np.array([2, 1, 0]) indices_blob = Const(read_init_net, idx, name="indices") reader = ds.random_reader(read_init_net, indices_blob, loop_over=True) reader.computeoffset(read_init_net) should_stop, batch = reader.read_record(read_next_net) workspace.CreateNet(read_init_net, True) workspace.RunNetOnce(read_init_net) workspace.CreateNet(read_next_net, True) for _ in range(len(entries) * 3): workspace.RunNet(str(read_next_net)) self.assertEquals(False, workspace.FetchBlob(should_stop)) """ 9. Sort and shuffle a dataset This sort the dataset using the score of a certain column, and then shuffle within each chunk of size batch_size * shuffle_size before shuffling the chunks. """ read_init_net = core.Net("read_init") read_next_net = core.Net("read_next") reader = ds.random_reader(read_init_net) reader.sort_and_shuffle(read_init_net, "int_lists:lengths", 1, 2) reader.computeoffset(read_init_net) should_continue, batch = reader.read_record(read_next_net) workspace.CreateNet(read_init_net, True) workspace.RunNetOnce(read_init_net) workspace.CreateNet(read_next_net, True) expected_idx = np.array([2, 1, 0]) for i in range(len(entries)): k = expected_idx[i] if i in expected_idx else i entry = entries[k] workspace.RunNet(str(read_next_net)) actual = FetchRecord(batch) _assert_records_equal(actual, entry) """ Trim a dataset """ trim_net = core.Net("trim_ds") ds.trim(trim_net, multiple_of=2) workspace.RunNetOnce(trim_net) trimmed = FetchRecord(ds.content()) EXPECTED_SIZES = [2, 2, 3, 3, 2, 2, 2, 6, 2, 3, 3, 4, 4, 2, 2, 2] actual_sizes = [d.shape[0] for d in trimmed.field_blobs()] self.assertEquals(EXPECTED_SIZES, actual_sizes) def test_last_n_window_ops(self): collect_net = core.Net("collect_net") collect_net.GivenTensorFill( [], "input", shape=[3, 2], values=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], ) input_array = np.array(list(range(1, 7)), dtype=np.float32).reshape(3, 2) workspace.CreateBlob("output") workspace.FeedBlob("next", np.array(0, dtype=np.int32)) collect_net.LastNWindowCollector( ["output", "next", "input"], ["output", "next"], num_to_collect=7, ) plan = core.Plan("collect_data") plan.AddStep(core.execution_step("collect_data", [collect_net], num_iter=1)) workspace.RunPlan(plan) reference_result = workspace.FetchBlob("output") npt.assert_array_equal(input_array, reference_result) plan = core.Plan("collect_data") plan.AddStep(core.execution_step("collect_data", [collect_net], num_iter=2)) workspace.RunPlan(plan) reference_result = workspace.FetchBlob("output") npt.assert_array_equal(input_array[[1, 2, 2, 0, 1, 2, 0]], reference_result) plan = core.Plan("collect_data") plan.AddStep(core.execution_step("collect_data", [collect_net], num_iter=3)) workspace.RunPlan(plan) reference_result = workspace.FetchBlob("output") npt.assert_array_equal(input_array[[2, 0, 1, 2, 2, 0, 1]], reference_result) def test_last_n_window_ops_shape_inference(self): collect_net = core.Net("collect_net") collect_net.GivenTensorFill( [], "input", shape=[3, 2], values=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], ) workspace.CreateBlob("output") workspace.FeedBlob("next", np.array(0, dtype=np.int32)) collect_net.LastNWindowCollector( ["output", "next", "input"], ["output", "next"], num_to_collect=7, ) (shapes, types) = workspace.InferShapesAndTypes([collect_net]) workspace.RunNetOnce(collect_net) self.assertTrue( np.array_equal( shapes["output"], np.array([7, workspace.blobs["output"].shape[1]]) ) ) def test_last_n_window_ops_shape_inference_4d_input(self): input_shape = [3, 2, 4, 5] collect_net = core.Net("collect_net") collect_net.GivenTensorFill( [], "input", shape=input_shape, values=[ float(val) for val in range(functools.reduce(operator.mul, input_shape)) ], ) workspace.CreateBlob("output") workspace.FeedBlob("next", np.array(0, dtype=np.int32)) collect_net.LastNWindowCollector( ["output", "next", "input"], ["output", "next"], num_to_collect=7, ) (shapes, types) = workspace.InferShapesAndTypes([collect_net]) workspace.RunNetOnce(collect_net) self.assertTrue( np.array_equal( shapes["output"], np.array([7, *list(workspace.blobs["output"].shape[1:])]) ) ) def test_collect_tensor_ops(self): init_net = core.Net("init_net") blobs = ["blob_1", "blob_2", "blob_3"] bvec_map = {} ONE = init_net.ConstantFill([], "ONE", shape=[1, 2], value=1) for b in blobs: init_net.ConstantFill([], [b], shape=[1, 2], value=0) bvec_map[b] = b + "_vec" init_net.CreateTensorVector([], [bvec_map[b]]) reader_net = core.Net("reader_net") for b in blobs: reader_net.Add([b, ONE], [b]) collect_net = core.Net("collect_net") num_to_collect = 1000 max_example_to_cover = 100000 bvec = [bvec_map[b] for b in blobs] collect_net.CollectTensor( bvec + blobs, bvec, num_to_collect=num_to_collect, ) print("Collect Net Proto: {}".format(collect_net.Proto())) plan = core.Plan("collect_data") plan.AddStep(core.execution_step("collect_init", init_net)) plan.AddStep( core.execution_step( "collect_data", [reader_net, collect_net], num_iter=max_example_to_cover ) ) workspace.RunPlan(plan) # concat the collected tensors concat_net = core.Net("concat_net") bconcated_map = {} bsize_map = {} for b in blobs: bconcated_map[b] = b + "_concated" bsize_map[b] = b + "_size" concat_net.ConcatTensorVector([bvec_map[b]], [bconcated_map[b]]) concat_net.TensorVectorSize([bvec_map[b]], [bsize_map[b]]) workspace.RunNetOnce(concat_net) # check data reference_result = workspace.FetchBlob(bconcated_map[blobs[0]]) self.assertEqual( reference_result.shape, (min(num_to_collect, max_example_to_cover), 2) ) size = workspace.FetchBlob(bsize_map[blobs[0]]) self.assertEqual(tuple(), size.shape) self.assertEqual(min(num_to_collect, max_example_to_cover), size.item()) hist, _ = np.histogram( reference_result[:, 0], bins=10, range=(1, max_example_to_cover) ) print("Sample histogram: {}".format(hist)) self.assertTrue(all(hist > 0.6 * (num_to_collect / 10))) for i in range(1, len(blobs)): result = workspace.FetchBlob(bconcated_map[blobs[i]]) self.assertEqual(reference_result.tolist(), result.tolist()) if __name__ == "__main__": import unittest unittest.main()