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