Dev flow.utils.data part1

oneflow/python/test/dataloader/test_numpy_dataset.py

@@ -0,0 +1,51 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import unittest
+import numpy as np
+
+import oneflow.experimental as flow
+import oneflow.python.utils.data as Data
+
+
+class ScpDataset(Data.Dataset):
+    def __init__(self, chunksize=200, dim=81, length=2000):
+        self.chunksize = chunksize
+        self.dim = dim
+        self.length = length
+
+    def __getitem__(self, index):
+        np.random.seed(index)
+        return np.random.randn(self.chunksize, self.dim)
+
+    def __len__(self):
+        return self.length
+
+
+@flow.unittest.skip_unless_1n1d()
+@unittest.skipIf(
+    not flow.unittest.env.eager_execution_enabled(),
+    ".numpy() doesn't work in lazy mode",
+)
+class TestNumpyDataset(flow.unittest.TestCase):
+    def test_numpy_dataset(test_case):
+        dataset = ScpDataset()
+        dataloader = Data.DataLoader(dataset, batch_size=32, shuffle=True)
+        for X in dataloader:
+            print(X.shape)
+
+
+if __name__ == "__main__":
+    unittest.main()

oneflow/python/test/dataloader/test_tensor_dataset.py

@@ -0,0 +1,82 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import unittest
+import numpy as np
+
+import oneflow.experimental as flow
+import oneflow.python.utils.data as Data
+import oneflow.experimental.nn as nn
+
+
+class LinearNet(nn.Module):
+    def __init__(self, n_feature):
+        super(LinearNet, self).__init__()
+        self.linear = nn.Linear(n_feature, 1)
+
+    def forward(self, x):
+        y = self.linear(x)
+        return y
+
+
+@flow.unittest.skip_unless_1n1d()
+@unittest.skipIf(
+    not flow.unittest.env.eager_execution_enabled(),
+    ".numpy() doesn't work in lazy mode",
+)
+class TestTensorDataset(flow.unittest.TestCase):
+    def test_tensor_dataset(test_case):
+
+        num_inputs = 2
+        num_examples = 1000
+        true_w = [2, -3.4]
+        true_b = 4.2
+
+        net = LinearNet(num_inputs)
+        print(net)
+        flow.nn.init.normal_(net.linear.weight, mean=0, std=0.01)
+        flow.nn.init.constant_(net.linear.bias, val=0)
+
+        loss = nn.MSELoss()
+        optimizer = flow.optim.SGD(net.parameters(), lr=0.03)
+        print(optimizer)
+
+        features = flow.tensor(
+            np.random.normal(0, 1, (num_examples, num_inputs)), dtype=flow.float
+        )
+        labels = true_w[0] * features[:, 0] + true_w[1] * features[:, 1] + true_b
+        labels += flow.tensor(
+            np.random.normal(0, 0.01, size=labels.size()), dtype=flow.float
+        )
+
+        batch_size = 10
+        # 将训练数据的特征和标签组合
+        dataset = Data.TensorDataset(features, labels)
+        # 随机读取小批量
+        data_iter = Data.DataLoader(dataset, batch_size, shuffle=True, num_workers=0)
+
+        num_epochs = 10
+        for epoch in range(1, num_epochs + 1):
+            for X, y in data_iter:
+                output = net(X)
+                l = loss(output, y)
+                optimizer.zero_grad()
+                l.backward()
+                optimizer.step()
+            print("epoch %d, loss: %f" % (epoch, l.numpy()))
+
+
+if __name__ == "__main__":
+    unittest.main()

oneflow/python/utils/data/__init__.py

@@ -0,0 +1,58 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+from oneflow.python.utils.data.sampler import (
+    Sampler,
+    SequentialSampler,
+    RandomSampler,
+    SubsetRandomSampler,
+    BatchSampler,
+)
+from oneflow.python.utils.data.dataset import (
+    Dataset,
+    IterableDataset,
+    TensorDataset,
+    ConcatDataset,
+    Subset,
+    random_split,
+)
+from oneflow.python.utils.data.dataset import IterableDataset as IterDataPipe
+from oneflow.python.utils.data.dataloader import DataLoader, _DatasetKind
+from oneflow.python.utils.data.decorator import (
+    functional_datapipe,
+    guaranteed_datapipes_determinism,
+    non_deterministic,
+)
+
+
+__all__ = [
+    "Sampler",
+    "SequentialSampler",
+    "RandomSampler",
+    "SubsetRandomSampler",
+    "BatchSampler",
+    "Dataset",
+    "IterableDataset",
+    "TensorDataset",
+    "ConcatDataset",
+    "Subset",
+    "random_split",
+    "DataLoader",
+    "_DatasetKind",
+    "IterDataPipe",
+    "functional_datapipe",
+    "guaranteed_datapipes_determinism",
+    "non_deterministic",
+]

oneflow/python/utils/data/_utils/__init__.py

@@ -0,0 +1,57 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+r"""Utility classes & functions for data loading. Code in this folder is mostly
+used by ../dataloder.py.
+
+A lot of multiprocessing is used in data loading, which only supports running
+functions defined in global environment (py2 can't serialize static methods).
+Therefore, for code tidiness we put these functions into different files in this
+folder.
+"""
+
+import sys
+import atexit
+
+
+IS_WINDOWS = sys.platform == "win32"
+
+MP_STATUS_CHECK_INTERVAL = 5.0
+r"""Interval (in seconds) to check status of processes to avoid hanging in
+    multiprocessing data loading. This is mainly used in getting data from
+    another process, in which case we need to periodically check whether the
+    sender is alive to prevent hanging."""
+
+
+python_exit_status = False
+r"""Whether Python is shutting down. This flag is guaranteed to be set before
+the Python core library resources are freed, but Python may already be exiting
+for some time when this is set.
+
+Hook to set this flag is `_set_python_exit_flag`, and is inspired by a similar
+hook in Python 3.7 multiprocessing library:
+https://github.com/python/cpython/blob/d4d60134b29290049e28df54f23493de4f1824b6/Lib/multiprocessing/util.py#L277-L327
+"""
+
+
+def _set_python_exit_flag():
+    global python_exit_status
+    python_exit_status = True
+
+
+atexit.register(_set_python_exit_flag)
+
+
+from . import collate, fetch

oneflow/python/utils/data/_utils/collate.py

@@ -0,0 +1,114 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+r""""Contains definitions of the methods used by the _BaseDataLoaderIter workers to
+collate samples fetched from dataset into Tensor(s).
+
+These **needs** to be in global scope since Py2 doesn't support serializing
+static methods.
+"""
+
+import oneflow as flow
+import re
+import collections
+import oneflow.python.utils as utils
+
+string_classes = (str, bytes)
+
+np_str_obj_array_pattern = re.compile(r"[SaUO]")
+
+
+def default_convert(data):
+    r"""Converts each NumPy array data field into a tensor"""
+    elem_type = type(data)
+    if isinstance(data, (flow.Tensor, flow._oneflow_internal.Tensor)):
+        return data
+    elif (
+        elem_type.__module__ == "numpy"
+        and elem_type.__name__ != "str_"
+        and elem_type.__name__ != "string_"
+    ):
+        # array of string classes and object
+        if (
+            elem_type.__name__ == "ndarray"
+            and np_str_obj_array_pattern.search(data.dtype.str) is not None
+        ):
+            return data
+        return flow.tensor(data)
+    elif isinstance(data, collections.abc.Mapping):
+        return {key: default_convert(data[key]) for key in data}
+    elif isinstance(data, tuple) and hasattr(data, "_fields"):  # namedtuple
+        return elem_type(*(default_convert(d) for d in data))
+    elif isinstance(data, collections.abc.Sequence) and not isinstance(
+        data, string_classes
+    ):
+        return [default_convert(d) for d in data]
+    else:
+        # NOTE: torch just return data here, and not raise any exception!
+        raise TypeError(default_convert_err_msg_format.format(elem_type))
+
+
+default_collate_err_msg_format = (
+    "default_collate: batch must contain tensors, numpy arrays, numbers, "
+    "dicts or lists; found {}"
+)
+
+default_convert_err_msg_format = (
+    "default_convert: batch must contain tensors, numpy arrays, numbers, "
+    "dicts or lists; found {}"
+)
+
+
+def default_collate(batch):
+    r"""Puts each data field into a tensor with outer dimension batch size"""
+
+    elem = batch[0]
+    elem_type = type(elem)
+    if isinstance(elem, (flow.Tensor, flow._oneflow_internal.Tensor)):
+        # TODO: tensor.storage()._new_shared(numel)
+        return flow.experimental.stack(batch, dim=0)
+    elif (
+        elem_type.__module__ == "numpy"
+        and elem_type.__name__ != "str_"
+        and elem_type.__name__ != "string_"
+    ):
+        if elem_type.__name__ == "ndarray" or elem_type.__name__ == "memmap":
+            # array of string classes and object
+            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
+                raise TypeError(default_collate_err_msg_format.format(elem.dtype))
+
+            return default_collate([flow.Tensor(b) for b in batch])
+        elif elem.shape == ():  # scalars
+            return flow.Tensor(batch)
+    elif isinstance(elem, float):
+        return flow.tensor(batch, dtype=flow.float64)
+    elif isinstance(elem, int):
+        return flow.tensor(batch)
+    elif isinstance(elem, string_classes):
+        return batch
+    elif isinstance(elem, collections.abc.Mapping):
+        return {key: default_collate([d[key] for d in batch]) for key in elem}
+    elif isinstance(elem, tuple) and hasattr(elem, "_fields"):  # namedtuple
+        return elem_type(*(default_collate(samples) for samples in zip(*batch)))
+    elif isinstance(elem, collections.abc.Sequence):
+        # check to make sure that the elements in batch have consistent size
+        it = iter(batch)
+        elem_size = len(next(it))
+        if not all(len(elem) == elem_size for elem in it):
+            raise RuntimeError("each element in list of batch should be of equal size")
+        transposed = zip(*batch)
+        return [default_collate(samples) for samples in transposed]
+
+    raise TypeError(default_collate_err_msg_format.format(elem_type))