default="Ascend",
choices=["Ascend", "GPU", "CPU"],
help="device target (default: Ascend)")
args = parser.parse_args()
if args.net == "squeezenet":
from src.squeezenet import SqueezeNet as squeezenet
else:
from src.squeezenet import SqueezeNet_Residual as squeezenet
if args.dataset == "cifar10":
num_classes = 10
else:
num_classes = 1000
context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
if args.device_target == "Ascend":
context.set_context(device_id=args.device_id)
if __name__ == '__main__':
net = squeezenet(num_classes=num_classes)
param_dict = load_checkpoint(args.ckpt_file)
load_param_into_net(net, param_dict)
input_data = Tensor(
np.zeros([args.batch_size, 3, args.height, args.width], np.float32))
export(net,
input_data,
file_name=args.file_name,
file_format=args.file_format)
# init context
device_id = os.getenv('DEVICE_ID')
device_id = int(device_id) if device_id else 0
context.set_context(mode=context.GRAPH_MODE,
device_target=target,
device_id=device_id)
# create dataset
dataset = create_dataset(dataset_path=args_opt.dataset_path,
do_train=False,
batch_size=config.batch_size,
target=target)
step_size = dataset.get_dataset_size()
# define net
net = squeezenet(num_classes=config.class_num)
# load checkpoint
param_dict = load_checkpoint(args_opt.checkpoint_path)
load_param_into_net(net, param_dict)
net.set_train(False)
# define loss
if args_opt.dataset == "imagenet":
if not config.use_label_smooth:
config.label_smooth_factor = 0.0
loss = CrossEntropySmooth(sparse=True,
reduction='mean',
smooth_factor=config.label_smooth_factor,
num_classes=config.class_num)
else: