forked from mindspore-Ecosystem/mindspore
googlenet-gpu
This commit is contained in:
parent
d4b5cda934
commit
7d5a67e9f0
|
@ -16,6 +16,8 @@
|
||||||
##############test googlenet example on cifar10#################
|
##############test googlenet example on cifar10#################
|
||||||
python eval.py
|
python eval.py
|
||||||
"""
|
"""
|
||||||
|
import argparse
|
||||||
|
|
||||||
import mindspore.nn as nn
|
import mindspore.nn as nn
|
||||||
from mindspore import context
|
from mindspore import context
|
||||||
from mindspore.nn.optim.momentum import Momentum
|
from mindspore.nn.optim.momentum import Momentum
|
||||||
|
@ -26,18 +28,27 @@ from src.config import cifar_cfg as cfg
|
||||||
from src.dataset import create_dataset
|
from src.dataset import create_dataset
|
||||||
from src.googlenet import GoogleNet
|
from src.googlenet import GoogleNet
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser(description='googlenet')
|
||||||
|
parser.add_argument('--checkpoint_path', type=str, default=None, help='Checkpoint file path')
|
||||||
|
args_opt = parser.parse_args()
|
||||||
|
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
|
device_target = cfg.device_target
|
||||||
context.set_context(mode=context.GRAPH_MODE, device_target=cfg.device_target)
|
context.set_context(mode=context.GRAPH_MODE, device_target=cfg.device_target)
|
||||||
context.set_context(device_id=cfg.device_id)
|
if device_target == "Ascend":
|
||||||
|
context.set_context(device_id=cfg.device_id)
|
||||||
|
|
||||||
net = GoogleNet(num_classes=cfg.num_classes)
|
net = GoogleNet(num_classes=cfg.num_classes, platform=device_target)
|
||||||
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, cfg.momentum,
|
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, cfg.momentum,
|
||||||
weight_decay=cfg.weight_decay)
|
weight_decay=cfg.weight_decay)
|
||||||
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean', is_grad=False)
|
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean', is_grad=False)
|
||||||
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})
|
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})
|
||||||
|
|
||||||
param_dict = load_checkpoint(cfg.checkpoint_path)
|
if device_target == "Ascend":
|
||||||
|
param_dict = load_checkpoint(cfg.checkpoint_path)
|
||||||
|
else: # GPU
|
||||||
|
param_dict = load_checkpoint(args_opt.checkpoint_path)
|
||||||
|
|
||||||
load_param_into_net(net, param_dict)
|
load_param_into_net(net, param_dict)
|
||||||
net.set_train(False)
|
net.set_train(False)
|
||||||
dataset = create_dataset(cfg.data_path, 1, False)
|
dataset = create_dataset(cfg.data_path, 1, False)
|
||||||
|
|
|
@ -0,0 +1,43 @@
|
||||||
|
#!/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.
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
ulimit -u unlimited
|
||||||
|
|
||||||
|
if [ $# != 1 ]
|
||||||
|
then
|
||||||
|
echo "GPU: sh run_eval_gpu.sh [CHECKPOINT_PATH]"
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
|
||||||
|
# check checkpoint file
|
||||||
|
if [ ! -f $1 ]
|
||||||
|
then
|
||||||
|
echo "error: CHECKPOINT_PATH=$1 is not a file"
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
|
||||||
|
BASEPATH=$(cd "`dirname $0`" || exit; pwd)
|
||||||
|
export PYTHONPATH=${BASEPATH}:$PYTHONPATH
|
||||||
|
export DEVICE_ID=0
|
||||||
|
|
||||||
|
if [ -d "../eval" ];
|
||||||
|
then
|
||||||
|
rm -rf ../eval
|
||||||
|
fi
|
||||||
|
mkdir ../eval
|
||||||
|
cd ../eval || exit
|
||||||
|
|
||||||
|
python3 ${BASEPATH}/../eval.py --checkpoint_path=$1 > ./eval.log 2>&1 &
|
|
@ -0,0 +1,45 @@
|
||||||
|
#!/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.
|
||||||
|
# ============================================================================
|
||||||
|
|
||||||
|
if [ $# -lt 2 ]
|
||||||
|
then
|
||||||
|
echo "Usage:\n \
|
||||||
|
sh run_train.sh [DEVICE_NUM] [VISIABLE_DEVICES(0,1,2,3,4,5,6,7)]\n \
|
||||||
|
"
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
|
||||||
|
if [ $1 -lt 1 ] && [ $1 -gt 8 ]
|
||||||
|
then
|
||||||
|
echo "error: DEVICE_NUM=$1 is not in (1-8)"
|
||||||
|
exit 1
|
||||||
|
fi
|
||||||
|
|
||||||
|
export DEVICE_NUM=$1
|
||||||
|
export RANK_SIZE=$1
|
||||||
|
|
||||||
|
BASEPATH=$(cd "`dirname $0`" || exit; pwd)
|
||||||
|
export PYTHONPATH=${BASEPATH}:$PYTHONPATH
|
||||||
|
if [ -d "../train" ];
|
||||||
|
then
|
||||||
|
rm -rf ../train
|
||||||
|
fi
|
||||||
|
mkdir ../train
|
||||||
|
cd ../train || exit
|
||||||
|
|
||||||
|
export CUDA_VISIBLE_DEVICES="$2"
|
||||||
|
mpirun -n $1 --allow-run-as-root \
|
||||||
|
python3 ${BASEPATH}/../train.py > train.log 2>&1 &
|
|
@ -56,24 +56,35 @@ class Inception(nn.Cell):
|
||||||
Inception Block
|
Inception Block
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, in_channels, n1x1, n3x3red, n3x3, n5x5red, n5x5, pool_planes):
|
def __init__(self, in_channels, n1x1, n3x3red, n3x3, n5x5red, n5x5, pool_planes, platform="Ascend"):
|
||||||
super(Inception, self).__init__()
|
super(Inception, self).__init__()
|
||||||
|
self.platform = platform
|
||||||
self.b1 = Conv2dBlock(in_channels, n1x1, kernel_size=1)
|
self.b1 = Conv2dBlock(in_channels, n1x1, kernel_size=1)
|
||||||
self.b2 = nn.SequentialCell([Conv2dBlock(in_channels, n3x3red, kernel_size=1),
|
self.b2 = nn.SequentialCell([Conv2dBlock(in_channels, n3x3red, kernel_size=1),
|
||||||
Conv2dBlock(n3x3red, n3x3, kernel_size=3, padding=0)])
|
Conv2dBlock(n3x3red, n3x3, kernel_size=3, padding=0)])
|
||||||
self.b3 = nn.SequentialCell([Conv2dBlock(in_channels, n5x5red, kernel_size=1),
|
self.b3 = nn.SequentialCell([Conv2dBlock(in_channels, n5x5red, kernel_size=1),
|
||||||
Conv2dBlock(n5x5red, n5x5, kernel_size=3, padding=0)])
|
Conv2dBlock(n5x5red, n5x5, kernel_size=3, padding=0)])
|
||||||
self.maxpool = P.MaxPoolWithArgmax(ksize=3, strides=1, padding="same")
|
if self.platform == "Ascend":
|
||||||
|
self.maxpool = P.MaxPoolWithArgmax(ksize=3, strides=1, padding="same")
|
||||||
|
else: # GPU
|
||||||
|
self.maxpool = P.MaxPool(ksize=3, strides=1, padding="same")
|
||||||
self.b4 = Conv2dBlock(in_channels, pool_planes, kernel_size=1)
|
self.b4 = Conv2dBlock(in_channels, pool_planes, kernel_size=1)
|
||||||
self.concat = P.Concat(axis=1)
|
self.concat = P.Concat(axis=1)
|
||||||
|
|
||||||
def construct(self, x):
|
def construct(self, x):
|
||||||
|
'''
|
||||||
|
construct inception model
|
||||||
|
'''
|
||||||
branch1 = self.b1(x)
|
branch1 = self.b1(x)
|
||||||
branch2 = self.b2(x)
|
branch2 = self.b2(x)
|
||||||
branch3 = self.b3(x)
|
branch3 = self.b3(x)
|
||||||
cell, argmax = self.maxpool(x)
|
if self.platform == "Ascend":
|
||||||
branch4 = self.b4(cell)
|
cell, argmax = self.maxpool(x)
|
||||||
_ = argmax
|
branch4 = self.b4(cell)
|
||||||
|
_ = argmax
|
||||||
|
else: # GPU
|
||||||
|
cell = self.maxpool(x)
|
||||||
|
branch4 = self.b4(cell)
|
||||||
return self.concat((branch1, branch2, branch3, branch4))
|
return self.concat((branch1, branch2, branch3, branch4))
|
||||||
|
|
||||||
|
|
||||||
|
@ -82,61 +93,82 @@ class GoogleNet(nn.Cell):
|
||||||
Googlenet architecture
|
Googlenet architecture
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, num_classes):
|
def __init__(self, num_classes, platform="Ascend"):
|
||||||
super(GoogleNet, self).__init__()
|
super(GoogleNet, self).__init__()
|
||||||
self.conv1 = Conv2dBlock(3, 64, kernel_size=7, stride=2, padding=0)
|
self.conv1 = Conv2dBlock(3, 64, kernel_size=7, stride=2, padding=0)
|
||||||
self.maxpool1 = P.MaxPoolWithArgmax(ksize=3, strides=2, padding="same")
|
self.platform = platform
|
||||||
|
|
||||||
self.conv2 = Conv2dBlock(64, 64, kernel_size=1)
|
self.conv2 = Conv2dBlock(64, 64, kernel_size=1)
|
||||||
self.conv3 = Conv2dBlock(64, 192, kernel_size=3, padding=0)
|
self.conv3 = Conv2dBlock(64, 192, kernel_size=3, padding=0)
|
||||||
self.maxpool2 = P.MaxPoolWithArgmax(ksize=3, strides=2, padding="same")
|
self.block3a = Inception(192, 64, 96, 128, 16, 32, 32, platform=self.platform)
|
||||||
|
self.block3b = Inception(256, 128, 128, 192, 32, 96, 64, platform=self.platform)
|
||||||
self.block3a = Inception(192, 64, 96, 128, 16, 32, 32)
|
self.block4a = Inception(480, 192, 96, 208, 16, 48, 64, platform=self.platform)
|
||||||
self.block3b = Inception(256, 128, 128, 192, 32, 96, 64)
|
self.block4b = Inception(512, 160, 112, 224, 24, 64, 64, platform=self.platform)
|
||||||
self.maxpool3 = P.MaxPoolWithArgmax(ksize=3, strides=2, padding="same")
|
self.block4c = Inception(512, 128, 128, 256, 24, 64, 64, platform=self.platform)
|
||||||
|
self.block4d = Inception(512, 112, 144, 288, 32, 64, 64, platform=self.platform)
|
||||||
self.block4a = Inception(480, 192, 96, 208, 16, 48, 64)
|
self.block4e = Inception(528, 256, 160, 320, 32, 128, 128, platform=self.platform)
|
||||||
self.block4b = Inception(512, 160, 112, 224, 24, 64, 64)
|
if self.platform == "Ascend":
|
||||||
self.block4c = Inception(512, 128, 128, 256, 24, 64, 64)
|
self.maxpool1 = P.MaxPoolWithArgmax(ksize=3, strides=2, padding="same")
|
||||||
self.block4d = Inception(512, 112, 144, 288, 32, 64, 64)
|
self.maxpool2 = P.MaxPoolWithArgmax(ksize=3, strides=2, padding="same")
|
||||||
self.block4e = Inception(528, 256, 160, 320, 32, 128, 128)
|
self.maxpool3 = P.MaxPoolWithArgmax(ksize=3, strides=2, padding="same")
|
||||||
self.maxpool4 = P.MaxPoolWithArgmax(ksize=2, strides=2, padding="same")
|
self.maxpool4 = P.MaxPoolWithArgmax(ksize=2, strides=2, padding="same")
|
||||||
|
else: # GPU
|
||||||
self.block5a = Inception(832, 256, 160, 320, 32, 128, 128)
|
self.maxpool1 = P.MaxPool(ksize=3, strides=2, padding="same")
|
||||||
self.block5b = Inception(832, 384, 192, 384, 48, 128, 128)
|
self.maxpool2 = P.MaxPool(ksize=3, strides=2, padding="same")
|
||||||
|
self.maxpool3 = P.MaxPool(ksize=3, strides=2, padding="same")
|
||||||
|
self.maxpool4 = P.MaxPool(ksize=2, strides=2, padding="same")
|
||||||
|
self.block5a = Inception(832, 256, 160, 320, 32, 128, 128, platform=self.platform)
|
||||||
|
self.block5b = Inception(832, 384, 192, 384, 48, 128, 128, platform=self.platform)
|
||||||
self.mean = P.ReduceMean(keep_dims=True)
|
self.mean = P.ReduceMean(keep_dims=True)
|
||||||
self.dropout = nn.Dropout(keep_prob=0.8)
|
self.dropout = nn.Dropout(keep_prob=0.8)
|
||||||
self.flatten = nn.Flatten()
|
self.flatten = nn.Flatten()
|
||||||
self.classifier = nn.Dense(1024, num_classes, weight_init=weight_variable(),
|
self.classifier = nn.Dense(1024, num_classes, weight_init=weight_variable(),
|
||||||
bias_init=weight_variable())
|
bias_init=weight_variable())
|
||||||
|
|
||||||
|
|
||||||
def construct(self, x):
|
def construct(self, x):
|
||||||
|
'''
|
||||||
|
construct googlenet model
|
||||||
|
'''
|
||||||
x = self.conv1(x)
|
x = self.conv1(x)
|
||||||
x, argmax = self.maxpool1(x)
|
if self.platform == "Ascend":
|
||||||
|
x, argmax = self.maxpool1(x)
|
||||||
|
else: # GPU
|
||||||
|
x = self.maxpool1(x)
|
||||||
|
|
||||||
x = self.conv2(x)
|
x = self.conv2(x)
|
||||||
x = self.conv3(x)
|
x = self.conv3(x)
|
||||||
x, argmax = self.maxpool2(x)
|
if self.platform == "Ascend":
|
||||||
|
x, argmax = self.maxpool2(x)
|
||||||
|
else: # GPU
|
||||||
|
x = self.maxpool2(x)
|
||||||
|
|
||||||
x = self.block3a(x)
|
x = self.block3a(x)
|
||||||
x = self.block3b(x)
|
x = self.block3b(x)
|
||||||
x, argmax = self.maxpool3(x)
|
if self.platform == "Ascend":
|
||||||
|
x, argmax = self.maxpool3(x)
|
||||||
|
else: # GPU
|
||||||
|
x = self.maxpool3(x)
|
||||||
|
|
||||||
x = self.block4a(x)
|
x = self.block4a(x)
|
||||||
x = self.block4b(x)
|
x = self.block4b(x)
|
||||||
x = self.block4c(x)
|
x = self.block4c(x)
|
||||||
x = self.block4d(x)
|
x = self.block4d(x)
|
||||||
x = self.block4e(x)
|
x = self.block4e(x)
|
||||||
x, argmax = self.maxpool4(x)
|
if self.platform == "Ascend":
|
||||||
|
x, argmax = self.maxpool4(x)
|
||||||
|
x = self.block5a(x)
|
||||||
|
x = self.block5b(x)
|
||||||
|
|
||||||
x = self.block5a(x)
|
x = self.mean(x, (2, 3))
|
||||||
x = self.block5b(x)
|
x = self.flatten(x)
|
||||||
|
x = self.classifier(x)
|
||||||
|
_ = argmax
|
||||||
|
else: # GPU
|
||||||
|
x = self.maxpool4(x)
|
||||||
|
x = self.block5a(x)
|
||||||
|
x = self.block5b(x)
|
||||||
|
|
||||||
x = self.mean(x, (2, 3))
|
x = self.mean(x, (2, 3))
|
||||||
x = self.flatten(x)
|
x = self.flatten(x)
|
||||||
x = self.classifier(x)
|
x = self.classifier(x)
|
||||||
|
|
||||||
_ = argmax
|
|
||||||
return x
|
return x
|
||||||
|
|
|
@ -25,7 +25,7 @@ import numpy as np
|
||||||
import mindspore.nn as nn
|
import mindspore.nn as nn
|
||||||
from mindspore import Tensor
|
from mindspore import Tensor
|
||||||
from mindspore import context
|
from mindspore import context
|
||||||
from mindspore.communication.management import init
|
from mindspore.communication.management import init, get_rank
|
||||||
from mindspore.nn.optim.momentum import Momentum
|
from mindspore.nn.optim.momentum import Momentum
|
||||||
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
|
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
|
||||||
from mindspore.train.model import Model, ParallelMode
|
from mindspore.train.model import Model, ParallelMode
|
||||||
|
@ -38,7 +38,6 @@ from src.googlenet import GoogleNet
|
||||||
random.seed(1)
|
random.seed(1)
|
||||||
np.random.seed(1)
|
np.random.seed(1)
|
||||||
|
|
||||||
|
|
||||||
def lr_steps(global_step, lr_max=None, total_epochs=None, steps_per_epoch=None):
|
def lr_steps(global_step, lr_max=None, total_epochs=None, steps_per_epoch=None):
|
||||||
"""Set learning rate."""
|
"""Set learning rate."""
|
||||||
lr_each_step = []
|
lr_each_step = []
|
||||||
|
@ -65,23 +64,36 @@ if __name__ == '__main__':
|
||||||
parser.add_argument('--device_id', type=int, default=None, help='device id of GPU or Ascend. (Default: None)')
|
parser.add_argument('--device_id', type=int, default=None, help='device id of GPU or Ascend. (Default: None)')
|
||||||
args_opt = parser.parse_args()
|
args_opt = parser.parse_args()
|
||||||
|
|
||||||
context.set_context(mode=context.GRAPH_MODE, device_target=cfg.device_target)
|
device_target = cfg.device_target
|
||||||
if args_opt.device_id is not None:
|
|
||||||
context.set_context(device_id=args_opt.device_id)
|
|
||||||
else:
|
|
||||||
context.set_context(device_id=cfg.device_id)
|
|
||||||
|
|
||||||
|
context.set_context(mode=context.GRAPH_MODE, device_target=cfg.device_target)
|
||||||
device_num = int(os.environ.get("DEVICE_NUM", 1))
|
device_num = int(os.environ.get("DEVICE_NUM", 1))
|
||||||
if device_num > 1:
|
|
||||||
context.reset_auto_parallel_context()
|
if device_target == "Ascend":
|
||||||
context.set_auto_parallel_context(device_num=device_num, parallel_mode=ParallelMode.DATA_PARALLEL,
|
if args_opt.device_id is not None:
|
||||||
mirror_mean=True)
|
context.set_context(device_id=args_opt.device_id)
|
||||||
init()
|
else:
|
||||||
|
context.set_context(device_id=cfg.device_id)
|
||||||
|
|
||||||
|
if device_num > 1:
|
||||||
|
context.reset_auto_parallel_context()
|
||||||
|
context.set_auto_parallel_context(device_num=device_num, parallel_mode=ParallelMode.DATA_PARALLEL,
|
||||||
|
mirror_mean=True)
|
||||||
|
init()
|
||||||
|
elif device_target == "GPU":
|
||||||
|
init("nccl")
|
||||||
|
|
||||||
|
if device_num > 1:
|
||||||
|
context.reset_auto_parallel_context()
|
||||||
|
context.set_auto_parallel_context(device_num=device_num, parallel_mode=ParallelMode.DATA_PARALLEL,
|
||||||
|
mirror_mean=True)
|
||||||
|
else:
|
||||||
|
raise ValueError("Unsupport platform.")
|
||||||
|
|
||||||
dataset = create_dataset(cfg.data_path, 1)
|
dataset = create_dataset(cfg.data_path, 1)
|
||||||
batch_num = dataset.get_dataset_size()
|
batch_num = dataset.get_dataset_size()
|
||||||
|
|
||||||
net = GoogleNet(num_classes=cfg.num_classes)
|
net = GoogleNet(num_classes=cfg.num_classes, platform=device_target)
|
||||||
# Continue training if set pre_trained to be True
|
# Continue training if set pre_trained to be True
|
||||||
if cfg.pre_trained:
|
if cfg.pre_trained:
|
||||||
param_dict = load_checkpoint(cfg.checkpoint_path)
|
param_dict = load_checkpoint(cfg.checkpoint_path)
|
||||||
|
@ -90,12 +102,19 @@ if __name__ == '__main__':
|
||||||
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), Tensor(lr), cfg.momentum,
|
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), Tensor(lr), cfg.momentum,
|
||||||
weight_decay=cfg.weight_decay)
|
weight_decay=cfg.weight_decay)
|
||||||
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean', is_grad=False)
|
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean', is_grad=False)
|
||||||
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'},
|
|
||||||
amp_level="O2", keep_batchnorm_fp32=False, loss_scale_manager=None)
|
if device_target == "Ascend":
|
||||||
|
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'},
|
||||||
|
amp_level="O2", keep_batchnorm_fp32=False, loss_scale_manager=None)
|
||||||
|
ckpt_save_dir = "./"
|
||||||
|
else: # GPU
|
||||||
|
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'},
|
||||||
|
amp_level="O2", keep_batchnorm_fp32=True, loss_scale_manager=None)
|
||||||
|
ckpt_save_dir = "./ckpt_" + str(get_rank()) + "/"
|
||||||
|
|
||||||
config_ck = CheckpointConfig(save_checkpoint_steps=batch_num * 5, keep_checkpoint_max=cfg.keep_checkpoint_max)
|
config_ck = CheckpointConfig(save_checkpoint_steps=batch_num * 5, keep_checkpoint_max=cfg.keep_checkpoint_max)
|
||||||
time_cb = TimeMonitor(data_size=batch_num)
|
time_cb = TimeMonitor(data_size=batch_num)
|
||||||
ckpoint_cb = ModelCheckpoint(prefix="train_googlenet_cifar10", directory="./", config=config_ck)
|
ckpoint_cb = ModelCheckpoint(prefix="train_googlenet_cifar10", directory=ckpt_save_dir, config=config_ck)
|
||||||
loss_cb = LossMonitor()
|
loss_cb = LossMonitor()
|
||||||
model.train(cfg.epoch_size, dataset, callbacks=[time_cb, ckpoint_cb, loss_cb])
|
model.train(cfg.epoch_size, dataset, callbacks=[time_cb, ckpoint_cb, loss_cb])
|
||||||
print("train success")
|
print("train success")
|
||||||
|
|
Loading…
Reference in New Issue