Browse Source

!10665 fix ST probability failure of resnet_thor in daily version

From: @wangmin0104
Reviewed-by: @wang_zi_dong,@sunnybeike
Signed-off-by: @sunnybeike
tags/v1.2.0-rc1
mindspore-ci-bot Gitee 5 years ago
parent
commit
f00ad3f755
1 changed files with 7 additions and 3 deletions
  1. +7
    -3
      tests/st/networks/models/resnet50/src/dataset.py

+ 7
- 3
tests/st/networks/models/resnet50/src/dataset.py View File

@@ -38,10 +38,14 @@ def create_dataset(dataset_path, do_train, repeat_num=1, batch_size=32):

device_num = int(os.getenv("RANK_SIZE"))
rank_id = int(os.getenv("RANK_ID"))
if device_num == 1:
data_set = ds.ImageFolderDataset(dataset_path, num_parallel_workers=8, shuffle=True)
if do_train:
if device_num == 1:
data_set = ds.ImageFolderDataset(dataset_path, num_parallel_workers=8, shuffle=True)
else:
data_set = ds.ImageFolderDataset(dataset_path, num_parallel_workers=8, shuffle=True,
num_shards=device_num, shard_id=rank_id)
else:
data_set = ds.ImageFolderDataset(dataset_path, num_parallel_workers=8, shuffle=True,
data_set = ds.ImageFolderDataset(dataset_path, num_parallel_workers=8, shuffle=False,
num_shards=device_num, shard_id=rank_id)

image_size = 224


Loading…
Cancel
Save