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)
batch_num = dataset.get_dataset_size()
net = GoogleNet(num_classes=cfg.num_classes)
# Continue training if set pre_trained to be True
if cfg.pre_trained:
param_dict = load_checkpoint(cfg.checkpoint_path)
load_param_into_net(net, param_dict)
lr = lr_steps(0,
lr_max=cfg.lr_init,
total_epochs=cfg.epoch_size,
steps_per_epoch=batch_num)
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()),
Tensor(lr),
cfg.momentum,
weight_decay=cfg.weight_decay)
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True,
reduction='mean',
is_grad=False)
def googlenet(*args, **kwargs):
return GoogleNet(*args, **kwargs)
from src.config import cifar_cfg, imagenet_cfg
from src.googlenet import GoogleNet
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Classification')
parser.add_argument('--dataset_name',
type=str,
default='cifar10',
choices=['imagenet', 'cifar10'],
help='dataset name.')
args_opt = parser.parse_args()
if args_opt.dataset_name == 'cifar10':
cfg = cifar_cfg
elif args_opt.dataset_name == 'imagenet':
cfg = imagenet_cfg
else:
raise ValueError("dataset is not support.")
net = GoogleNet(num_classes=cfg.num_classes)
assert cfg.checkpoint_path is not None, "cfg.checkpoint_path is None."
param_dict = load_checkpoint(cfg.checkpoint_path)
load_param_into_net(net, param_dict)
input_arr = Tensor(np.random.uniform(0.0, 1.0, size=[1, 3, 224, 224]),
ms.float32)
export(net, input_arr, file_name=cfg.onnx_filename, file_format="ONNX")
export(net, input_arr, file_name=cfg.air_filename, file_format="AIR")
from src.config import cifar_cfg as cfg
from src.dataset import create_dataset
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__':
device_target = cfg.device_target
context.set_context(mode=context.GRAPH_MODE, device_target=cfg.device_target)
if device_target == "Ascend":
context.set_context(device_id=cfg.device_id)
net = GoogleNet(num_classes=cfg.num_classes)
opt = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, cfg.momentum,
weight_decay=cfg.weight_decay)
loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean', is_grad=False)
model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc'})
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)
net.set_train(False)
dataset = create_dataset(cfg.data_path, 1, False)
acc = model.eval(dataset)
print("accuracy: ", acc)