Voyage
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix tinydarknet

This commit is contained in:
l_emon 2021-07-10 17:32:43 +08:00
parent bcaafbf68d
commit 12b78e2a9b
1 changed files with 10 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
model_zoo/official/cv/tinydarknet/train.py
View File

@ -120,17 +120,6 @@ def modelarts_pre_process():
@moxing_wrapper(pre_process=modelarts_pre_process)
def run_train():
if config.dataset_name == "imagenet":
dataset = create_dataset_imagenet(config.train_data_dir, 1)
elif config.dataset_name == "cifar10":
dataset = create_dataset_cifar(dataset_path=config.train_data_dir,
do_train=True,
repeat_num=1,
batch_size=config.batch_size,
target=config.device_target)
else:
raise ValueError("Unsupported dataset.")
# set context
device_target = config.device_target
@ -148,7 +137,16 @@ def run_train():
gradients_mean=True)
init()
rank = get_rank_id()
if config.dataset_name == "imagenet":
dataset = create_dataset_imagenet(config.train_data_dir, 1)
elif config.dataset_name == "cifar10":
dataset = create_dataset_cifar(dataset_path=config.train_data_dir,
do_train=True,
repeat_num=1,
batch_size=config.batch_size,
target=config.device_target)
else:
raise ValueError("Unsupported dataset.")
batch_num = dataset.get_dataset_size()
net = TinyDarkNet(num_classes=config.num_classes)