| |
|
| |
|
| |
|
| |
|
| |
|
| | import unittest |
| | import numpy as np |
| |
|
| | import caffe2.proto.caffe2_pb2 as caffe2_pb2 |
| | from caffe2.python import core, workspace, timeout_guard |
| |
|
| |
|
| | @unittest.skipIf(not workspace.C.use_mkldnn, "No MKLDNN support.") |
| | class BlobsQueueDBTest(unittest.TestCase): |
| | def test_create_blobs_queue_db_string(self): |
| | device_opt = core.DeviceOption(caffe2_pb2.IDEEP, 0) |
| | with core.DeviceScope(device_opt): |
| | def add_blobs(queue, num_samples): |
| | blob = core.BlobReference("blob") |
| | status = core.BlobReference("blob_status") |
| | for i in range(num_samples): |
| | self._add_blob_to_queue( |
| | queue, self._create_test_tensor_protos(i), blob, status |
| | ) |
| | self._test_create_blobs_queue_db(add_blobs) |
| |
|
| | def test_create_blobs_queue_db_tensor(self): |
| | device_opt = core.DeviceOption(caffe2_pb2.IDEEP, 0) |
| | with core.DeviceScope(device_opt): |
| | def add_blobs(queue, num_samples): |
| | blob = core.BlobReference("blob") |
| | status = core.BlobReference("blob_status") |
| | for i in range(num_samples): |
| | data = self._create_test_tensor_protos(i) |
| | data = np.array([data], dtype=str) |
| | self._add_blob_to_queue( |
| | queue, data, blob, status |
| | ) |
| | self._test_create_blobs_queue_db(add_blobs) |
| |
|
| | def _test_create_blobs_queue_db(self, add_blobs_fun): |
| | device_opt = core.DeviceOption(caffe2_pb2.IDEEP, 0) |
| | with core.DeviceScope(device_opt): |
| | num_samples = 10000 |
| | batch_size = 10 |
| | init_net = core.Net('init_net') |
| | net = core.Net('test_create_blobs_queue_db') |
| | queue = init_net.CreateBlobsQueue([], 'queue', capacity=num_samples) |
| | reader = init_net.CreateBlobsQueueDB( |
| | [queue], |
| | 'blobs_queue_db_reader', |
| | value_blob_index=0, |
| | timeout_secs=0.1, |
| | ) |
| | workspace.RunNetOnce(init_net) |
| | add_blobs_fun(queue, num_samples) |
| |
|
| | net.TensorProtosDBInput( |
| | [reader], |
| | ['image', 'label'], |
| | batch_size=batch_size |
| | ) |
| | workspace.CreateNet(net) |
| |
|
| | close_net = core.Net('close_net') |
| | close_net.CloseBlobsQueue([queue], []) |
| |
|
| | for i in range(int(num_samples / batch_size)): |
| | with timeout_guard.CompleteInTimeOrDie(2.0): |
| | workspace.RunNet(net) |
| |
|
| | images = workspace.FetchBlob('image') |
| | labels = workspace.FetchBlob('label') |
| | self.assertEqual(batch_size, len(images)) |
| | self.assertEqual(batch_size, len(labels)) |
| | for idx, item in enumerate(images): |
| | self.assertEqual( |
| | "foo{}".format(i * batch_size + idx).encode('utf-8'), item |
| | ) |
| | for item in labels: |
| | self.assertEqual(1, item) |
| | workspace.RunNetOnce(close_net) |
| |
|
| | def _add_blob_to_queue(self, queue, data, blob, status): |
| | device_opt = core.DeviceOption(caffe2_pb2.IDEEP, 0) |
| | with core.DeviceScope(device_opt): |
| | workspace.FeedBlob(blob, data, core.DeviceOption(caffe2_pb2.CPU, 0)) |
| | op = core.CreateOperator( |
| | "SafeEnqueueBlobs", |
| | [queue, blob], |
| | [blob, status], |
| | ) |
| |
|
| | workspace.RunOperatorOnce(op) |
| |
|
| | def _create_test_tensor_protos(self, idx): |
| | item = caffe2_pb2.TensorProtos() |
| | data = item.protos.add() |
| | data.data_type = core.DataType.STRING |
| | data.string_data.append("foo{}".format(idx).encode('utf-8')) |
| | label = item.protos.add() |
| | label.data_type = core.DataType.INT32 |
| | label.int32_data.append(1) |
| |
|
| | return item.SerializeToString() |
| |
|
| | if __name__ == "__main__": |
| | import unittest |
| | unittest.main() |
| |
|
