1# Copyright 2020 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================ 15"""dataset base.""" 16import os 17 18from mindspore import dataset as ds 19from mindspore.common import dtype as mstype 20from mindspore.dataset.transforms import c_transforms as C 21from mindspore.dataset.vision import Inter 22from mindspore.dataset.vision import c_transforms as CV 23 24 25def create_mnist_dataset(mode='train', num_samples=2, batch_size=2): 26 """create dataset for train or test""" 27 mnist_path = '/home/workspace/mindspore_dataset/mnist' 28 num_parallel_workers = 1 29 30 # define dataset 31 mnist_ds = ds.MnistDataset(os.path.join(mnist_path, mode), num_samples=num_samples, shuffle=False) 32 33 resize_height, resize_width = 32, 32 34 35 # define map operations 36 resize_op = CV.Resize((resize_height, resize_width), interpolation=Inter.LINEAR) # Bilinear mode 37 rescale_nml_op = CV.Rescale(1 / 0.3081, -1 * 0.1307 / 0.3081) 38 rescale_op = CV.Rescale(1.0 / 255.0, shift=0.0) 39 hwc2chw_op = CV.HWC2CHW() 40 type_cast_op = C.TypeCast(mstype.int32) 41 42 # apply map operations on images 43 mnist_ds = mnist_ds.map(operations=type_cast_op, input_columns="label", num_parallel_workers=num_parallel_workers) 44 mnist_ds = mnist_ds.map(operations=resize_op, input_columns="image", num_parallel_workers=num_parallel_workers) 45 mnist_ds = mnist_ds.map(operations=rescale_op, input_columns="image", num_parallel_workers=num_parallel_workers) 46 mnist_ds = mnist_ds.map(operations=rescale_nml_op, input_columns="image", num_parallel_workers=num_parallel_workers) 47 mnist_ds = mnist_ds.map(operations=hwc2chw_op, input_columns="image", num_parallel_workers=num_parallel_workers) 48 49 # apply DatasetOps 50 mnist_ds = mnist_ds.batch(batch_size=batch_size, drop_remainder=True) 51 52 return mnist_ds 53