diff --git a/cnn_e2e/ofrecord_util.py b/cnn_e2e/ofrecord_util.py index f72efa7d20104475fe87e625c4a570ce493df7ec..6a2d9ca15a86ded1e60a18ec95e1657aa64ba837 100644 --- a/cnn_e2e/ofrecord_util.py +++ b/cnn_e2e/ofrecord_util.py @@ -6,6 +6,7 @@ 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") @@ -88,22 +89,21 @@ 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 - seed = 0 color_space = 'RGB' - with flow.fixed_placement("cpu", "0:0"): + with flow.fixed_placement("cpu", "0:0-{}".format(args.gpu_num_per_node - 1)): 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, 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) - rsz = flow.image.Resize(image, resize_x=args.image_size, - resize_y=args.image_size, + 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) return label, normal @@ -112,23 +112,19 @@ 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 - seed = 0 color_space = 'RGB' - with flow.fixed_placement("cpu", "0:0"): + with flow.fixed_placement("cpu", "0:0-{}".format(args.gpu_num_per_node - 1)): 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.OFRecordImageDecoderRandomCrop(ofrecord, "encoded", seed=seed, - 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_x=args.image_size, - resize_y=args.image_size, - color_space=color_space) + rsz = flow.image.Resize(image, resize_shorter=args.resize_shorter, color_space=color_space) - rng = flow.random.CoinFlip(batch_size=val_batch_size, seed=seed) - normal = flow.image.CropMirrorNormalize(rsz, mirror_blob=rng, 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) return label, normal