support the MNIST-series datasets

Author: lzhengning

LICENSE.txt

@@ -50,6 +50,11 @@ def inner(*args, **kw):
 from . import autograd
 from .autograd import *
 
+def ndimension(self):
+    return self.ndim
+Var.ndimension = ndimension
+
+
 Tensor = Var
 tensor = wrapper(array)
 
@@ -76,3 +81,17 @@ def forward(self, *args, **kw):
 import jtorch.nn
 from jtorch.nn import Module, Parameter
 import jtorch.optim
+
+from jtorch.utils.dtype import Dtype, get_string_dtype
+
+def frombuffer(buffer: bytearray, 
+              *, 
+              dtype: Dtype, 
+              count: int = -1, 
+              offset: int = 0, 
+              requires_grad: bool = True) -> Tensor:
+    dtype = get_string_dtype(dtype)
+    tensor = jt.array(np.frombuffer(buffer, dtype, count=count, offset=offset))
+    if requires_grad and tensor.dtype.is_float():
+        tensor.requires_grad = True
+    return tensor

python/jtorch/__init__.py

@@ -19,4 +19,40 @@
     train=False,
     download=True,
     transform=ToTensor(),
-)
+)
+
+batch_size = 64
+
+# Create data loaders.
+train_dataloader = DataLoader(training_data, batch_size=batch_size)
+test_dataloader = DataLoader(test_data, batch_size=batch_size)
+
+for X, y in test_dataloader:
+    print(f"Shape of X [N, C, H, W]: {X.shape}")
+    print(f"Shape of y: {y.shape} {y.dtype}")
+    break
+
+# Get cpu or gpu device for training.
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using {device} device")
+
+# Define model
+class NeuralNetwork(nn.Module):
+    def __init__(self):
+        super(NeuralNetwork, self).__init__()
+        self.flatten = nn.Flatten()
+        self.linear_relu_stack = nn.Sequential(
+            nn.Linear(28*28, 512),
+            nn.ReLU(),
+            nn.Linear(512, 512),
+            nn.ReLU(),
+            nn.Linear(512, 10)
+        )
+
+    def forward(self, x):
+        x = self.flatten(x)
+        logits = self.linear_relu_stack(x)
+        return logits
+
+model = NeuralNetwork().to(device)
+print(model)

python/jtorch/tutorial/quickstart.py

@@ -1,2 +1,56 @@
-class DataLoader:
-    pass
+from audioop import mul
+import jittor as jt
+from jittor.dataset import Dataset as JDataset
+
+from typing import Any, Callable, Iterable, Optional, Sequence, Union
+
+
+class Dataset:
+    def __getitem__(self, index):
+        raise NotImplementedError
+
+
+class DataLoader(JDataset):
+    def __init__(self, dataset: Dataset, 
+                 batch_size: Optional[int] = 1,
+                 shuffle: Optional[bool] = None, 
+                 sampler = None,
+                 batch_sampler = None,
+                 num_workers: int = 0, 
+                 collate_fn = None,
+                 pin_memory: bool = False, 
+                 drop_last: bool = False,
+                 timeout: float = 0, 
+                 worker_init_fn = None,
+                 multiprocessing_context=None, 
+                 generator=None,
+                 *, prefetch_factor: int = 2,
+                 persistent_workers: bool = False,
+                 pin_memory_device: str = "") -> None:
+        super().__init__(batch_size=batch_size, 
+                         shuffle=shuffle,
+                         num_workers=num_workers,
+                         drop_last=drop_last)
+        
+        unsupported_kwargs = {
+            "sampler": sampler, 
+            "batch_sampler": batch_sampler, 
+            "pin_memory": pin_memory, 
+            "timeout": timeout,
+            "worker_init_fn": worker_init_fn,
+            "multiprocessing_context": multiprocessing_context, 
+            "generator": generator, 
+            "persistent_workers": persistent_workers, 
+            "pin_memory_device": pin_memory_device
+        }
+        for kwarg, value in unsupported_kwargs.items():
+            if value:
+                jt.LOG.w(f"Not implemented Dataloader kwarg: {kwarg}")
+
+        self.collate_fn = collate_fn
+        
+    def collate_batch(self, batch):
+        if self.collate_fn is not None:
+            return self.collate_fn(batch)
+        else:
+            return super().collate_batch(batch)

python/jtorch/vision/datasets.py renamed to python/jtorch/utils/__init__.py

@@ -0,0 +1,9 @@
+from typing import Callable, Union
+Dtype = Union[Callable, str]
+
+def get_string_dtype(dtype):
+    if callable(dtype):
+        dtype = dtype.__name__
+    if not isinstance(dtype, str):
+        raise ValueError(f"dtype is expected to be str, python type function, or jittor type function, but got {dtype}.")
+    return dtype

python/jtorch/utils/data.py

@@ -0,0 +1,46 @@
+import importlib.machinery
+import os
+
+
+def _download_file_from_remote_location(fpath: str, url: str) -> None:
+    pass
+
+
+def _is_remote_location_available() -> bool:
+    return False
+
+
+def _get_extension_path(lib_name):
+
+    lib_dir = os.path.dirname(__file__)
+    if os.name == "nt":
+        # Register the main torchvision library location on the default DLL path
+        import ctypes
+        import sys
+
+        kernel32 = ctypes.WinDLL("kernel32.dll", use_last_error=True)
+        with_load_library_flags = hasattr(kernel32, "AddDllDirectory")
+        prev_error_mode = kernel32.SetErrorMode(0x0001)
+
+        if with_load_library_flags:
+            kernel32.AddDllDirectory.restype = ctypes.c_void_p
+
+        if sys.version_info >= (3, 8):
+            os.add_dll_directory(lib_dir)
+        elif with_load_library_flags:
+            res = kernel32.AddDllDirectory(lib_dir)
+            if res is None:
+                err = ctypes.WinError(ctypes.get_last_error())
+                err.strerror += f' Error adding "{lib_dir}" to the DLL directories.'
+                raise err
+
+        kernel32.SetErrorMode(prev_error_mode)
+
+    loader_details = (importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES)
+
+    extfinder = importlib.machinery.FileFinder(lib_dir, loader_details)
+    ext_specs = extfinder.find_spec(lib_name)
+    if ext_specs is None:
+        raise ImportError
+
+    return ext_specs.origin

python/jtorch/utils/dtype.py

@@ -0,0 +1,9 @@
+from .mnist import EMNIST, FashionMNIST, KMNIST, MNIST, QMNIST
+
+__all__ = (
+    "EMNIST",
+    "FashionMNIST",
+    "QMNIST",
+    "MNIST",
+    "KMNIST",
+)