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 16 17import mindspore.dataset as ds 18import mindspore.dataset.transforms.c_transforms as ops 19 20 21def test_random_choice(): 22 """ 23 Test RandomChoice op 24 """ 25 ds.config.set_seed(0) 26 27 def test_config(arr, op_list): 28 try: 29 data = ds.NumpySlicesDataset(arr, column_names="col", shuffle=False) 30 data = data.map(operations=ops.RandomChoice(op_list), input_columns=["col"]) 31 res = [] 32 for i in data.create_dict_iterator(num_epochs=1, output_numpy=True): 33 res.append(i["col"].tolist()) 34 return res 35 except (TypeError, ValueError) as e: 36 return str(e) 37 38 # Test whether an operation would be randomly chosen. 39 # In order to prevent random failure, both results need to be checked. 40 res1 = test_config([[0, 1, 2]], [ops.PadEnd([4], 0), ops.Slice([0, 2])]) 41 assert res1 in [[[0, 1, 2, 0]], [[0, 2]]] 42 43 # Test nested structure 44 res2 = test_config([[0, 1, 2]], [ops.Compose([ops.Duplicate(), ops.Concatenate()]), 45 ops.Compose([ops.Slice([0, 1]), ops.OneHot(2)])]) 46 assert res2 in [[[[1, 0], [0, 1]]], [[0, 1, 2, 0, 1, 2]]] 47 # Test RandomChoice when there is only 1 operation 48 assert test_config([[4, 3], [2, 1]], [ops.Slice([0])]) == [[4], [2]] 49 50 51if __name__ == "__main__": 52 test_random_choice() 53