1:modify shell for deeplabv3

2:fix normalize bug
3:add ci test3:add ci test3:add ci test
This commit is contained in:
unknown 2020-06-09 11:38:07 +08:00
parent 09433cdd06
commit c929dab3dd
11 changed files with 201 additions and 55 deletions

View File

@ -16,17 +16,17 @@ This is an example of training DeepLabv3 with PASCAL VOC 2012 dataset in MindSpo
- Set options in config.py. - Set options in config.py.
- Run `run_standalone_train.sh` for non-distributed training. - Run `run_standalone_train.sh` for non-distributed training.
``` bash ``` bash
sh scripts/run_standalone_train.sh DEVICE_ID EPOCH_SIZE DATA_DIR sh scripts/run_standalone_train.sh DEVICE_ID DATA_PATH
``` ```
- Run `run_distribute_train.sh` for distributed training. - Run `run_distribute_train.sh` for distributed training.
``` bash ``` bash
sh scripts/run_distribute_train.sh DEVICE_NUM EPOCH_SIZE DATA_DIR MINDSPORE_HCCL_CONFIG_PATH sh scripts/run_distribute_train.sh MINDSPORE_HCCL_CONFIG_PATH DATA_PATH
``` ```
### Evaluation ### Evaluation
Set options in evaluation_config.py. Make sure the 'data_file' and 'finetune_ckpt' are set to your own path. Set options in evaluation_config.py. Make sure the 'data_file' and 'finetune_ckpt' are set to your own path.
- Run run_eval.sh for evaluation. - Run run_eval.sh for evaluation.
``` bash ``` bash
sh scripts/run_eval.sh DEVICE_ID DATA_DIR sh scripts/run_eval.sh DEVICE_ID DATA_PATH PRETRAINED_CKPT_PATH
``` ```
## Options and Parameters ## Options and Parameters
@ -49,6 +49,11 @@ config.py:
decoder_output_stride The ratio of input to output spatial resolution when employing decoder decoder_output_stride The ratio of input to output spatial resolution when employing decoder
to refine segmentation results, default is None. to refine segmentation results, default is None.
image_pyramid Input scales for multi-scale feature extraction, default is None. image_pyramid Input scales for multi-scale feature extraction, default is None.
epoch_size Epoch size, default is 6.
batch_size batch size of input dataset: N, default is 2.
enable_save_ckpt Enable save checkpoint, default is true.
save_checkpoint_steps Save checkpoint steps, default is 1000.
save_checkpoint_num Save checkpoint numbers, default is 1.
``` ```
@ -56,11 +61,6 @@ config.py:
``` ```
Parameters for dataset and network: Parameters for dataset and network:
distribute Run distribute, default is false. distribute Run distribute, default is false.
epoch_size Epoch size, default is 6.
batch_size batch size of input dataset: N, default is 2.
data_url Train/Evaluation data url, required. data_url Train/Evaluation data url, required.
checkpoint_url Checkpoint path, default is None. checkpoint_url Checkpoint path, default is None.
enable_save_ckpt Enable save checkpoint, default is true.
save_checkpoint_steps Save checkpoint steps, default is 1000.
save_checkpoint_num Save checkpoint numbers, default is 1.
``` ```

View File

@ -25,9 +25,7 @@ from src.config import config
parser = argparse.ArgumentParser(description="Deeplabv3 evaluation") parser = argparse.ArgumentParser(description="Deeplabv3 evaluation")
parser.add_argument('--epoch_size', type=int, default=2, help='Epoch size.')
parser.add_argument("--device_id", type=int, default=0, help="Device id, default is 0.") parser.add_argument("--device_id", type=int, default=0, help="Device id, default is 0.")
parser.add_argument('--batch_size', type=int, default=2, help='Batch size.')
parser.add_argument('--data_url', required=True, default=None, help='Evaluation data url') parser.add_argument('--data_url', required=True, default=None, help='Evaluation data url')
parser.add_argument('--checkpoint_url', default=None, help='Checkpoint path') parser.add_argument('--checkpoint_url', default=None, help='Checkpoint path')
@ -39,8 +37,8 @@ print(args_opt)
if __name__ == "__main__": if __name__ == "__main__":
args_opt.crop_size = config.crop_size args_opt.crop_size = config.crop_size
args_opt.base_size = config.crop_size args_opt.base_size = config.crop_size
eval_dataset = create_dataset(args_opt, args_opt.data_url, args_opt.epoch_size, args_opt.batch_size, usage="eval") eval_dataset = create_dataset(args_opt, args_opt.data_url, config.epoch_size, config.batch_size, usage="eval")
net = deeplabv3_resnet50(config.seg_num_classes, [args_opt.batch_size, 3, args_opt.crop_size, args_opt.crop_size], net = deeplabv3_resnet50(config.seg_num_classes, [config.batch_size, 3, args_opt.crop_size, args_opt.crop_size],
infer_scale_sizes=config.eval_scales, atrous_rates=config.atrous_rates, infer_scale_sizes=config.eval_scales, atrous_rates=config.atrous_rates,
decoder_output_stride=config.decoder_output_stride, output_stride=config.output_stride, decoder_output_stride=config.decoder_output_stride, output_stride=config.output_stride,
fine_tune_batch_norm=config.fine_tune_batch_norm, image_pyramid=config.image_pyramid) fine_tune_batch_norm=config.fine_tune_batch_norm, image_pyramid=config.image_pyramid)

View File

@ -16,17 +16,21 @@
echo "==============================================================================================================" echo "=============================================================================================================="
echo "Please run the scipt as: " echo "Please run the scipt as: "
echo "bash run_distribute_train.sh DEVICE_NUM EPOCH_SIZE DATA_DIR MINDSPORE_HCCL_CONFIG_PATH" echo "bash run_distribute_train.sh MINDSPORE_HCCL_CONFIG_PATH DATA_PATH"
echo "for example: bash run_distribute_train.sh 8 40 /path/zh-wiki/ /path/hccl.json" echo "for example: bash run_distribute_train.sh MINDSPORE_HCCL_CONFIG_PATH DATA_PATH [PRETRAINED_CKPT_PATH](option)"
echo "It is better to use absolute path." echo "It is better to use absolute path."
echo "==============================================================================================================" echo "=============================================================================================================="
EPOCH_SIZE=$2 DATA_DIR=$2
DATA_DIR=$3
export MINDSPORE_HCCL_CONFIG_PATH=$4 export MINDSPORE_HCCL_CONFIG_PATH=$1
export RANK_TABLE_FILE=$4 export RANK_TABLE_FILE=$1
export RANK_SIZE=$1 export RANK_SIZE=8
PATH_CHECKPOINT=""
if [ $# == 3 ]
then
PATH_CHECKPOINT=$3
fi
cores=`cat /proc/cpuinfo|grep "processor" |wc -l` cores=`cat /proc/cpuinfo|grep "processor" |wc -l`
echo "the number of logical core" $cores echo "the number of logical core" $cores
avg_core_per_rank=`expr $cores \/ $RANK_SIZE` avg_core_per_rank=`expr $cores \/ $RANK_SIZE`
@ -55,12 +59,8 @@ do
env > env.log env > env.log
taskset -c $cmdopt python ../train.py \ taskset -c $cmdopt python ../train.py \
--distribute="true" \ --distribute="true" \
--epoch_size=$EPOCH_SIZE \
--device_id=$DEVICE_ID \ --device_id=$DEVICE_ID \
--enable_save_ckpt="true" \ --checkpoint_url=$PATH_CHECKPOINT \
--checkpoint_url="" \
--save_checkpoint_steps=10000 \
--save_checkpoint_num=1 \
--data_url=$DATA_DIR > log.txt 2>&1 & --data_url=$DATA_DIR > log.txt 2>&1 &
cd ../ cd ../
done done

View File

@ -15,18 +15,20 @@
# ============================================================================ # ============================================================================
echo "==============================================================================================================" echo "=============================================================================================================="
echo "Please run the scipt as: " echo "Please run the scipt as: "
echo "bash run_eval.sh DEVICE_ID DATA_DIR" echo "bash run_eval.sh DEVICE_ID DATA_PATH PRETRAINED_CKPT_PATH"
echo "for example: bash run_eval.sh /path/zh-wiki/ " echo "for example: bash run_eval.sh DEVICE_ID DATA_PATH PRETRAINED_CKPT_PATH"
echo "==============================================================================================================" echo "=============================================================================================================="
DEVICE_ID=$1 DEVICE_ID=$1
DATA_DIR=$2 DATA_DIR=$2
PATH_CHECKPOINT=$3
mkdir -p ms_log mkdir -p ms_log
CUR_DIR=`pwd` CUR_DIR=`pwd`
export GLOG_log_dir=${CUR_DIR}/ms_log export GLOG_log_dir=${CUR_DIR}/ms_log
export GLOG_logtostderr=0 export GLOG_logtostderr=0
python evaluation.py \ python eval.py \
--device_id=$DEVICE_ID \ --device_id=$DEVICE_ID \
--checkpoint_url="" \ --checkpoint_url=$PATH_CHECKPOINT \
--data_url=$DATA_DIR > log.txt 2>&1 & --data_url=$DATA_DIR > log.txt 2>&1 &

View File

@ -15,13 +15,17 @@
# ============================================================================ # ============================================================================
echo "==============================================================================================================" echo "=============================================================================================================="
echo "Please run the scipt as: " echo "Please run the scipt as: "
echo "bash run_standalone_pretrain.sh DEVICE_ID EPOCH_SIZE DATA_DIR" echo "bash run_standalone_pretrain.sh DEVICE_ID DATA_PATH"
echo "for example: bash run_standalone_train.sh 0 40 /path/zh-wiki/ " echo "for example: bash run_standalone_train.sh DEVICE_ID DATA_PATH [PRETRAINED_CKPT_PATH](option)"
echo "==============================================================================================================" echo "=============================================================================================================="
DEVICE_ID=$1 DEVICE_ID=$1
EPOCH_SIZE=$2 DATA_DIR=$2
DATA_DIR=$3 PATH_CHECKPOINT=""
if [ $# == 3 ]
then
PATH_CHECKPOINT=$3
fi
mkdir -p ms_log mkdir -p ms_log
CUR_DIR=`pwd` CUR_DIR=`pwd`
@ -29,10 +33,6 @@ export GLOG_log_dir=${CUR_DIR}/ms_log
export GLOG_logtostderr=0 export GLOG_logtostderr=0
python train.py \ python train.py \
--distribute="false" \ --distribute="false" \
--epoch_size=$EPOCH_SIZE \
--device_id=$DEVICE_ID \ --device_id=$DEVICE_ID \
--enable_save_ckpt="true" \ --checkpoint_url=$PATH_CHECKPOINT \
--checkpoint_url="" \
--save_checkpoint_steps=10000 \
--save_checkpoint_num=1 \
--data_url=$DATA_DIR > log.txt 2>&1 & --data_url=$DATA_DIR > log.txt 2>&1 &

View File

@ -29,5 +29,10 @@ config = ed({
"fine_tune_batch_norm": False, "fine_tune_batch_norm": False,
"ignore_label": 255, "ignore_label": 255,
"decoder_output_stride": None, "decoder_output_stride": None,
"seg_num_classes": 21 "seg_num_classes": 21,
"epoch_size": 6,
"batch_size": 2,
"enable_save_ckpt": True,
"save_checkpoint_steps": 10000,
"save_checkpoint_num": 1
}) })

View File

@ -16,6 +16,7 @@
from PIL import Image from PIL import Image
import mindspore.dataset as de import mindspore.dataset as de
import mindspore.dataset.transforms.vision.c_transforms as C import mindspore.dataset.transforms.vision.c_transforms as C
import numpy as np
from .ei_dataset import HwVocRawDataset from .ei_dataset import HwVocRawDataset
from .utils import custom_transforms as tr from .utils import custom_transforms as tr
@ -52,8 +53,8 @@ class DataTransform:
rhf_tr = tr.RandomHorizontalFlip() rhf_tr = tr.RandomHorizontalFlip()
image, label = rhf_tr(image, label) image, label = rhf_tr(image, label)
nor_tr = tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) image = np.array(image).astype(np.float32)
image, label = nor_tr(image, label) label = np.array(label).astype(np.float32)
return image, label return image, label
@ -71,13 +72,13 @@ class DataTransform:
fsc_tr = tr.FixScaleCrop(crop_size=self.args.crop_size) fsc_tr = tr.FixScaleCrop(crop_size=self.args.crop_size)
image, label = fsc_tr(image, label) image, label = fsc_tr(image, label)
nor_tr = tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) image = np.array(image).astype(np.float32)
image, label = nor_tr(image, label) label = np.array(label).astype(np.float32)
return image, label return image, label
def create_dataset(args, data_url, epoch_num=1, batch_size=1, usage="train"): def create_dataset(args, data_url, epoch_num=1, batch_size=1, usage="train", shuffle=True):
""" """
Create Dataset for DeepLabV3. Create Dataset for DeepLabV3.
@ -106,7 +107,7 @@ def create_dataset(args, data_url, epoch_num=1, batch_size=1, usage="train"):
# 1464 samples / batch_size 8 = 183 batches # 1464 samples / batch_size 8 = 183 batches
# epoch_num is num of steps # epoch_num is num of steps
# 3658 steps / 183 = 20 epochs # 3658 steps / 183 = 20 epochs
if usage == "train": if usage == "train" and shuffle:
dataset = dataset.shuffle(1464) dataset = dataset.shuffle(1464)
dataset = dataset.batch(batch_size, drop_remainder=(usage == "train")) dataset = dataset.batch(batch_size, drop_remainder=(usage == "train"))
dataset = dataset.repeat(count=epoch_num) dataset = dataset.repeat(count=epoch_num)

View File

@ -33,6 +33,7 @@ class Normalize:
def __call__(self, img, mask): def __call__(self, img, mask):
img = np.array(img).astype(np.float32) img = np.array(img).astype(np.float32)
mask = np.array(mask).astype(np.float32) mask = np.array(mask).astype(np.float32)
img = ((img - self.mean) / self.std).astype(np.float32)
return img, mask return img, mask

View File

@ -27,14 +27,10 @@ from src.config import config
parser = argparse.ArgumentParser(description="Deeplabv3 training") parser = argparse.ArgumentParser(description="Deeplabv3 training")
parser.add_argument("--distribute", type=str, default="false", help="Run distribute, default is false.") parser.add_argument("--distribute", type=str, default="false", help="Run distribute, default is false.")
parser.add_argument('--epoch_size', type=int, default=6, help='Epoch size.')
parser.add_argument('--batch_size', type=int, default=2, help='Batch size.')
parser.add_argument('--data_url', required=True, default=None, help='Train data url') parser.add_argument('--data_url', required=True, default=None, help='Train data url')
parser.add_argument("--device_id", type=int, default=0, help="Device id, default is 0.") parser.add_argument("--device_id", type=int, default=0, help="Device id, default is 0.")
parser.add_argument('--checkpoint_url', default=None, help='Checkpoint path') parser.add_argument('--checkpoint_url', default=None, help='Checkpoint path')
parser.add_argument("--enable_save_ckpt", type=str, default="true", help="Enable save checkpoint, default is true.")
parser.add_argument("--save_checkpoint_steps", type=int, default=1000, help="Save checkpoint steps, default is 1000.")
parser.add_argument("--save_checkpoint_num", type=int, default=1, help="Save checkpoint numbers, default is 1.")
args_opt = parser.parse_args() args_opt = parser.parse_args()
print(args_opt) print(args_opt)
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args_opt.device_id) context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args_opt.device_id)
@ -70,16 +66,16 @@ if __name__ == "__main__":
init() init()
args_opt.base_size = config.crop_size args_opt.base_size = config.crop_size
args_opt.crop_size = config.crop_size args_opt.crop_size = config.crop_size
train_dataset = create_dataset(args_opt, args_opt.data_url, args_opt.epoch_size, args_opt.batch_size, usage="train") train_dataset = create_dataset(args_opt, args_opt.data_url, config.epoch_size, config.batch_size, usage="train")
dataset_size = train_dataset.get_dataset_size() dataset_size = train_dataset.get_dataset_size()
time_cb = TimeMonitor(data_size=dataset_size) time_cb = TimeMonitor(data_size=dataset_size)
callback = [time_cb, LossCallBack()] callback = [time_cb, LossCallBack()]
if args_opt.enable_save_ckpt == "true": if config.enable_save_ckpt:
config_ck = CheckpointConfig(save_checkpoint_steps=args_opt.save_checkpoint_steps, config_ck = CheckpointConfig(save_checkpoint_steps=config.save_checkpoint_steps,
keep_checkpoint_max=args_opt.save_checkpoint_num) keep_checkpoint_max=config.save_checkpoint_num)
ckpoint_cb = ModelCheckpoint(prefix='checkpoint_deeplabv3', config=config_ck) ckpoint_cb = ModelCheckpoint(prefix='checkpoint_deeplabv3', config=config_ck)
callback.append(ckpoint_cb) callback.append(ckpoint_cb)
net = deeplabv3_resnet50(config.seg_num_classes, [args_opt.batch_size, 3, args_opt.crop_size, args_opt.crop_size], net = deeplabv3_resnet50(config.seg_num_classes, [config.batch_size, 3, args_opt.crop_size, args_opt.crop_size],
infer_scale_sizes=config.eval_scales, atrous_rates=config.atrous_rates, infer_scale_sizes=config.eval_scales, atrous_rates=config.atrous_rates,
decoder_output_stride=config.decoder_output_stride, output_stride=config.output_stride, decoder_output_stride=config.decoder_output_stride, output_stride=config.output_stride,
fine_tune_batch_norm=config.fine_tune_batch_norm, image_pyramid=config.image_pyramid) fine_tune_batch_norm=config.fine_tune_batch_norm, image_pyramid=config.image_pyramid)
@ -88,5 +84,5 @@ if __name__ == "__main__":
loss = OhemLoss(config.seg_num_classes, config.ignore_label) loss = OhemLoss(config.seg_num_classes, config.ignore_label)
opt = Momentum(filter(lambda x: 'beta' not in x.name and 'gamma' not in x.name and 'depth' not in x.name and 'bias' not in x.name, net.trainable_params()), learning_rate=config.learning_rate, momentum=config.momentum, weight_decay=config.weight_decay) opt = Momentum(filter(lambda x: 'beta' not in x.name and 'gamma' not in x.name and 'depth' not in x.name and 'bias' not in x.name, net.trainable_params()), learning_rate=config.learning_rate, momentum=config.momentum, weight_decay=config.weight_decay)
model = Model(net, loss, opt) model = Model(net, loss, opt)
model.train(args_opt.epoch_size, train_dataset, callback) model.train(config.epoch_size, train_dataset, callback)

View File

@ -0,0 +1,47 @@
#!/bin/bash
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
echo "=============================================================================================================="
echo "Please run the scipt as: "
echo "for example: bash run_deeplabv3_ci.sh DEVICE_ID DATA_PATH PRETRAINED_CKPT_PATH"
echo "=============================================================================================================="
DEVICE_ID=$1
DATA_DIR=$2
PATH_CHECKPOINT=$3
BASE_PATH=$(cd "$(dirname $0)"; pwd)
unset SLOG_PRINT_TO_STDOUT
CODE_DIR="./"
if [ -d ${BASE_PATH}/../../../../model_zoo/deeplabv3 ]; then
CODE_DIR=${BASE_PATH}/../../../../model_zoo/deeplabv3
elif [ -d ${BASE_PATH}/../../model_zoo/deeplabv3 ]; then
CODE_DIR=${BASE_PATH}/../../model_zoo/deeplabv3
else
echo "[ERROR] code dir is not found"
fi
echo $CODE_DIR
rm -rf ${BASE_PATH}/deeplabv3
cp -r ${CODE_DIR} ${BASE_PATH}/deeplabv3
cp -f ${BASE_PATH}/train_one_epoch_with_loss.py ${BASE_PATH}/deeplabv3/train_one_epoch_with_loss.py
cd ${BASE_PATH}/deeplabv3
python train_one_epoch_with_loss.py --data_url=$DATA_DIR --checkpoint_url=$PATH_CHECKPOINT --device_id=$DEVICE_ID > train_deeplabv3_ci.log 2>&1 &
process_pid=`echo $!`
wait ${process_pid}
status=`echo $?`
if [ "${status}" != "0" ]; then
echo "[ERROR] test deeplabv3 failed. status: ${status}"
exit 1
else
echo "[INFO] test deeplabv3 success."
fi

View File

@ -0,0 +1,96 @@
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""train."""
import argparse
import time
from mindspore import context
from mindspore.nn.optim.momentum import Momentum
from mindspore import Model
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore.train.callback import Callback
from src.md_dataset import create_dataset
from src.losses import OhemLoss
from src.deeplabv3 import deeplabv3_resnet50
from src.config import config
parser = argparse.ArgumentParser(description="Deeplabv3 training")
parser.add_argument('--data_url', required=True, default=None, help='Train data url')
parser.add_argument("--device_id", type=int, default=0, help="Device id, default is 0.")
parser.add_argument('--checkpoint_url', default=None, help='Checkpoint path')
args_opt = parser.parse_args()
print(args_opt)
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=args_opt.device_id)
class LossCallBack(Callback):
"""
Monitor the loss in training.
Note:
if per_print_times is 0 do not print loss.
Args:
per_print_times (int): Print loss every times. Default: 1.
"""
def __init__(self, data_size, per_print_times=1):
super(LossCallBack, self).__init__()
if not isinstance(per_print_times, int) or per_print_times < 0:
raise ValueError("print_step must be int and >= 0")
self.data_size = data_size
self._per_print_times = per_print_times
self.time = 1000
self.loss = 0
def epoch_begin(self, run_context):
self.epoch_time = time.time()
def step_end(self, run_context):
cb_params = run_context.original_args()
epoch_mseconds = (time.time() - self.epoch_time) * 1000
self.time = epoch_mseconds / self.data_size
self.loss += cb_params.net_outputs
print("epoch: {}, step: {}, outputs are {}".format(cb_params.cur_epoch_num, cb_params.cur_step_num,
str(cb_params.net_outputs)))
def model_fine_tune(flags, train_net, fix_weight_layer):
checkpoint_path = flags.checkpoint_url
if checkpoint_path is None:
return
param_dict = load_checkpoint(checkpoint_path)
load_param_into_net(train_net, param_dict)
for para in train_net.trainable_params():
if fix_weight_layer in para.name:
para.requires_grad = False
if __name__ == "__main__":
start_time = time.time()
epoch_size = 3
args_opt.base_size = config.crop_size
args_opt.crop_size = config.crop_size
train_dataset = create_dataset(args_opt, args_opt.data_url, epoch_size, config.batch_size,
usage="train", shuffle=False)
dataset_size = train_dataset.get_dataset_size()
callback = LossCallBack(dataset_size)
net = deeplabv3_resnet50(config.seg_num_classes, [config.batch_size, 3, args_opt.crop_size, args_opt.crop_size],
infer_scale_sizes=config.eval_scales, atrous_rates=config.atrous_rates,
decoder_output_stride=config.decoder_output_stride, output_stride=config.output_stride,
fine_tune_batch_norm=config.fine_tune_batch_norm, image_pyramid=config.image_pyramid)
net.set_train()
model_fine_tune(args_opt, net, 'layer')
loss = OhemLoss(config.seg_num_classes, config.ignore_label)
opt = Momentum(filter(lambda x: 'beta' not in x.name and 'gamma' not in x.name and 'depth' not in x.name and 'bias' not in x.name, net.trainable_params()), learning_rate=config.learning_rate, momentum=config.momentum, weight_decay=config.weight_decay)
model = Model(net, loss, opt)
model.train(epoch_size, train_dataset, callback)
print(time.time() - start_time)
print("expect loss: ", callback.loss / 3)
print("expect time: ", callback.time)
expect_loss = 0.5
expect_time = 35
assert callback.loss.asnumpy() / 3 <= expect_loss
assert callback.time <= expect_time