return label, normal
if __name__ == "__main__":
import os
import config as configs
from util import Summary, InitNodes, Metric
from job_function_util import get_val_config
parser = configs.get_parser()
args = parser.parse_args()
configs.print_args(args)
flow.config.gpu_device_num(args.gpu_num_per_node)
flow.config.enable_debug_mode(True)
@flow.global_function(get_val_config(args))
def IOTest():
if args.train_data_dir:
assert os.path.exists(args.train_data_dir)
print("Loading data from {}".format(args.train_data_dir))
(labels, images) = load_imagenet_for_training(args)
else:
print("Loading synthetic data.")
(labels, images) = load_synthetic(args)
outputs = {"images": images, "labels": labels}
return outputs
total_device_num = args.num_nodes * args.gpu_num_per_node
train_batch_size = total_device_num * args.batch_size_per_device
summary = Summary(args.log_dir, args, filename='io_test.csv')
metric = Metric(desc='io_test',
one_hot_labels = label_smoothing(labels, args.num_classes, args.label_smoothing, logits.dtype)
loss = flow.nn.softmax_cross_entropy_with_logits(one_hot_labels, logits, name="softmax_loss")
else:
loss = flow.nn.sparse_softmax_cross_entropy_with_logits(labels, logits, name="softmax_loss")
if not args.use_fp16:
loss = flow.math.reduce_mean(loss)
predictions = flow.nn.softmax(logits)
outputs = {"loss": loss, "predictions": predictions, "labels": labels}
# set up warmup,learning rate and optimizer
optimizer_util.set_up_optimizer(loss, args)
return outputs
@flow.global_function("predict", get_val_config(args))
def InferenceNet():
if args.val_data_dir:
assert os.path.exists(args.val_data_dir)
print("Loading data from {}".format(args.val_data_dir))
(labels, images) = ofrecord_util.load_imagenet_for_validation(args)
else:
print("Loading synthetic data.")
(labels, images) = ofrecord_util.load_synthetic(args)
logits = model_dict[args.model](images, args)
predictions = flow.nn.softmax(logits)
outputs = {"predictions": predictions, "labels": labels}
return outputs
