diff --git a/cnn_benchmark/ofrecord_util.py b/cnn_benchmark/ofrecord_util.py index 6a2d9ca15a86ded1e60a18ec95e1657aa64ba837..cae67369103dd0b06d5eae1997a973e5fe9c84a4 100755 --- a/cnn_benchmark/ofrecord_util.py +++ b/cnn_benchmark/ofrecord_util.py @@ -4,13 +4,20 @@ from __future__ import print_function import oneflow as flow + def add_ofrecord_args(parser): - parser.add_argument("--image_size", type=int, default=224, required=False, help="image size") - parser.add_argument("--resize_shorter", type=int, default=256, required=False, help="resize shorter for validation") - parser.add_argument("--train_data_dir", type=str, default=None, help="train dataset directory") - parser.add_argument("--train_data_part_num", type=int, default=256, help="train data part num") - parser.add_argument("--val_data_dir", type=str, default=None, help="val dataset directory") - parser.add_argument("--val_data_part_num", type=int, default=256, help="val data part num") + parser.add_argument("--image_size", type=int, default=224, + required=False, help="image size") + parser.add_argument("--resize_shorter", type=int, default=256, + required=False, help="resize shorter for validation") + parser.add_argument("--train_data_dir", type=str, + default=None, help="train dataset directory") + parser.add_argument("--train_data_part_num", type=int, + default=256, help="train data part num") + parser.add_argument("--val_data_dir", type=str, + default=None, help="val dataset directory") + parser.add_argument("--val_data_part_num", type=int, + default=256, help="val data part num") return parser @@ -22,7 +29,7 @@ def load_imagenet(args, batch_size, data_dir, data_part_num, codec): codec=codec, preprocessors=[flow.data.NormByChannelPreprocessor(args.rgb_mean[::-1], args.rgb_std[::-1])], - #preprocessors=[flow.data.NormByChannelPreprocessor(args.rgb_mean, args.rgb_std)], #bgr2rgb + # preprocessors=[flow.data.NormByChannelPreprocessor(args.rgb_mean, args.rgb_std)], #bgr2rgb ) label_blob_conf = flow.data.BlobConf( @@ -36,7 +43,7 @@ def load_imagenet(args, batch_size, data_dir, data_part_num, codec): data_part_num=data_part_num, part_name_suffix_length=5, #shuffle = True, - #buffer_size=32768, + # buffer_size=32768, name="decode", ) @@ -44,8 +51,8 @@ def load_imagenet(args, batch_size, data_dir, data_part_num, codec): def load_imagenet_for_training(args): total_device_num = args.num_nodes * args.gpu_num_per_node train_batch_size = total_device_num * args.batch_size_per_device - codec=flow.data.ImageCodec([ - #flow.data.ImagePreprocessor('bgr2rgb'), + codec = flow.data.ImageCodec([ + # flow.data.ImagePreprocessor('bgr2rgb'), #flow.data.ImageCropWithRandomSizePreprocessor(area=(0.08, 1)), flow.data.ImageResizePreprocessor(args.image_size, args.image_size), flow.data.ImagePreprocessor('mirror'), @@ -57,12 +64,13 @@ def load_imagenet_for_training(args): def load_imagenet_for_validation(args): total_device_num = args.num_nodes * args.gpu_num_per_node val_batch_size = total_device_num * args.val_batch_size_per_device - codec=flow.data.ImageCodec( + codec = flow.data.ImageCodec( [ - #flow.data.ImagePreprocessor('bgr2rgb'), + # flow.data.ImagePreprocessor('bgr2rgb'), # flow.data.ImageTargetResizePreprocessor(resize_shorter=256), # flow.data.ImageCenterCropPreprocessor(args.image_size, args.image_size), - flow.data.ImageResizePreprocessor(args.image_size, args.image_size), + flow.data.ImageResizePreprocessor( + args.image_size, args.image_size), ] ) return load_imagenet(args, val_batch_size, args.val_data_dir, args.val_data_part_num, codec) @@ -85,49 +93,68 @@ def load_synthetic(args): return label, image +def get_placement(args): + + machine_node_list = [] + nodes_str = "0-{}".format(args.gpu_num_per_node - 1) + + for i in range(args.num_nodes): + machine_node_list.append("{}:{}".format(i, nodes_str)) + + return machine_node_list + + def load_imagenet_for_training2(args): total_device_num = args.num_nodes * args.gpu_num_per_node train_batch_size = total_device_num * args.batch_size_per_device color_space = 'RGB' - with flow.fixed_placement("cpu", "0:0-{}".format(args.gpu_num_per_node - 1)): + # with flow.fixed_placement("cpu", ["0:0-1", "1:0-1"]): + with flow.fixed_placement("cpu", get_placement(args)): ofrecord = flow.data.ofrecord_reader(args.train_data_dir, batch_size=train_batch_size, data_part_num=args.train_data_part_num, part_name_suffix_length=5, - random_shuffle = True, + random_shuffle=True, shuffle_after_epoch=True) - image = flow.data.OFRecordImageDecoderRandomCrop(ofrecord, "encoded", #seed=seed, + image = flow.data.OFRecordImageDecoderRandomCrop(ofrecord, "encoded", # seed=seed, color_space=color_space) - label = flow.data.OFRecordRawDecoder(ofrecord, "class/label", shape=(), dtype=flow.int32) + label = flow.data.OFRecordRawDecoder( + ofrecord, "class/label", shape=(), dtype=flow.int32) rsz = flow.image.Resize(image, resize_x=args.image_size, resize_y=args.image_size, color_space=color_space) - rng = flow.random.CoinFlip(batch_size=train_batch_size)#, seed=seed) + rng = flow.random.CoinFlip(batch_size=train_batch_size) # , seed=seed) normal = flow.image.CropMirrorNormalize(rsz, mirror_blob=rng, color_space=color_space, - mean=args.rgb_mean, std=args.rgb_std, output_dtype = flow.float) + mean=args.rgb_mean, std=args.rgb_std, output_dtype=flow.float) return label, normal + def load_imagenet_for_validation2(args): total_device_num = args.num_nodes * args.gpu_num_per_node val_batch_size = total_device_num * args.val_batch_size_per_device color_space = 'RGB' - with flow.fixed_placement("cpu", "0:0-{}".format(args.gpu_num_per_node - 1)): + # with flow.fixed_placement("cpu", ["0:0-1", "1:0-1"]): + with flow.fixed_placement("cpu", get_placement(args)): ofrecord = flow.data.ofrecord_reader(args.val_data_dir, batch_size=val_batch_size, data_part_num=args.val_data_part_num, part_name_suffix_length=5, shuffle_after_epoch=False) - image = flow.data.OFRecordImageDecoder(ofrecord, "encoded", color_space=color_space) - label = flow.data.OFRecordRawDecoder(ofrecord, "class/label", shape=(), dtype=flow.int32) - rsz = flow.image.Resize(image, resize_shorter=args.resize_shorter, color_space=color_space) + image = flow.data.OFRecordImageDecoder( + ofrecord, "encoded", color_space=color_space) + label = flow.data.OFRecordRawDecoder( + ofrecord, "class/label", shape=(), dtype=flow.int32) + rsz = flow.image.Resize( + image, resize_shorter=args.resize_shorter, color_space=color_space) normal = flow.image.CropMirrorNormalize(rsz, color_space=color_space, - crop_h = args.image_size, crop_w = args.image_size, crop_pos_y = 0.5, crop_pos_x = 0.5, - mean=args.rgb_mean, std=args.rgb_std, output_dtype = flow.float) + crop_h=args.image_size, crop_w=args.image_size, crop_pos_y=0.5, crop_pos_x=0.5, + mean=args.rgb_mean, std=args.rgb_std, output_dtype=flow.float) return label, normal + if __name__ == "__main__": import os import config as configs @@ -149,7 +176,7 @@ if __name__ == "__main__": else: print("Loading synthetic data.") (labels, images) = load_synthetic(args) - outputs = {"images":images, "labels": labels} + outputs = {"images": images, "labels": labels} return outputs total_device_num = args.num_nodes * args.gpu_num_per_node