ArantxaCasanova

First model version

a00ee36 almost 5 years ago

12.9 kB

# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
# All contributions by Andy Brock:
# Copyright (c) 2019 Andy Brock
#
# MIT License
""" Convert dataset to HDF5
    This script preprocesses a dataset and saves it (images and labels) to 
    an HDF5 file for improved I/O. """
import os
import sys
from argparse import ArgumentParser
from tqdm import tqdm, trange
import h5py as h5

import numpy as np
import torch

import utils
import torchvision.transforms.functional as tv_f


def prepare_parser():
    usage = "Parser for ImageNet HDF5 scripts."
    parser = ArgumentParser(description=usage)
    parser.add_argument(
        "--resolution",
        type=int,
        default=128,
        help="Which Dataset resolution to train on, out of 64, 128, 256 (default: %(default)s)",
    )
    parser.add_argument(
        "--split",
        type=str,
        default="train",
        help="Which Dataset to convert: train, val (default: %(default)s)",
    )
    parser.add_argument(
        "--data_root",
        type=str,
        default="data",
        help="Default location where data is stored (default: %(default)s)",
    )
    parser.add_argument(
        "--out_path",
        type=str,
        default="data",
        help="Default location where data in hdf5 format will be stored (default: %(default)s)",
    )
    parser.add_argument(
        "--pretrained_model_path",
        type=str,
        default="",
        help="Location where the pretrained model (to extract features) can be found (default: %(default)s)",
    )
    parser.add_argument(
        "--save_features_only",
        action="store_true",
        default=False,
        help="Only save features in hdf5 file.",
    )
    parser.add_argument(
        "--save_images_only",
        action="store_true",
        default=False,
        help="Only save images and their labels in hdf5 file.",
    )
    parser.add_argument(
        "--feature_augmentation",
        action="store_true",
        default=False,
        help="Additioally store instance features with horizontally flipped input images.",
    )
    parser.add_argument(
        "--feature_extractor",
        type=str,
        default="classification",
        choices=["classification", "selfsupervised"],
        help="Choice of feature extractor",
    )
    parser.add_argument(
        "--backbone_feature_extractor",
        type=str,
        default="resnet50",
        choices=["resnet50"],
        help="Choice of feature extractor backbone",
    )
    parser.add_argument(
        "--which_dataset", type=str, default="imagenet", help="Dataset choice."
    )
    parser.add_argument(
        "--instance_json",
        type=str,
        default="",
        help="Path to JSON containing instance segmentations for COCO_Stuff",
    )
    parser.add_argument(
        "--stuff_json",
        type=str,
        default="",
        help="Path to JSON containing instance segmentations for COCO_Stuff",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=256,
        help="Default overall batchsize (default: %(default)s)",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=16,
        help="Number of dataloader workers (default: %(default)s)",
    )
    parser.add_argument(
        "--chunk_size",
        type=int,
        default=500,
        help="Default overall batchsize (default: %(default)s)",
    )
    parser.add_argument(
        "--compression",
        action="store_true",
        default=False,
        help="Use LZF compression? (default: %(default)s)",
    )
    return parser


def run(config):
    # Load pretrained feature extractor
    net = utils.load_pretrained_feature_extractor(
        config["pretrained_model_path"],
        config["feature_extractor"],
        config["backbone_feature_extractor"],
    )
    net.eval()

    # Update compression entry
    config["compression"] = (
        "lzf" if config["compression"] else None
    )  # No compression; can also use 'lzf'

    # Get dataset
    kwargs = {
        "num_workers": config["num_workers"],
        "pin_memory": False,
        "drop_last": False,
    }
    test_part = False
    if config["split"] == "test":
        config["split"] = "val"
        test_part = True
    if config["which_dataset"] in ["imagenet", "imagenet_lt"]:
        data_path = os.path.join(config["data_root"], config["split"])
    else:
        data_path = config["data_root"]
    dataset = utils.get_dataset_images(
        config["resolution"],
        data_path=data_path,
        longtail=config["which_dataset"] == "imagenet_lt",
        split=config["split"],
        test_part=test_part,
        which_dataset=config["which_dataset"],
        instance_json=config["instance_json"],
        stuff_json=config["stuff_json"],
    )
    train_loader = utils.get_dataloader(
        dataset, config["batch_size"], shuffle=False, **kwargs
    )

    # HDF5 supports chunking and compression. You may want to experiment
    # with different chunk sizes to see how it runs on your machines.
    # Chunk Size/compression     Read speed @ 256x256   Read speed @ 128x128  Filesize @ 128x128    Time to write @128x128
    # 1 / None                   20/s
    # 500 / None                 ramps up to 77/s       102/s                 61GB                  23min
    # 500 / LZF                                         8/s                   56GB                  23min
    # 1000 / None                78/s
    # 5000 / None                81/s
    # auto:(125,1,16,32) / None                         11/s                  61GB

    # Use imagenet statistics to preprocess images for the feature extractor (instance features)
    norm_mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda()
    norm_std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda()

    if config["which_dataset"] in ["imagenet", "imagenet_lt"]:
        dataset_name_prefix = "ILSVRC"
    elif config["which_dataset"] == "coco":
        dataset_name_prefix = "COCO"
    else:
        dataset_name_prefix = config["which_dataset"]

    if not config["save_features_only"]:
        h5file_name = config["out_path"] + "/%s%i%s%s%s_xy.hdf5" % (
            dataset_name_prefix,
            config["resolution"],
            "" if config["which_dataset"] != "imagenet_lt" else "longtail",
            "_val" if config["split"] == "val" else "",
            "_test" if test_part else "",
        )
        print("Filenames are ", h5file_name)

    if not config["save_images_only"]:
        h5file_name_feats = config["out_path"] + "/%s%i%s%s%s_feats_%s_%s.hdf5" % (
            dataset_name_prefix,
            config["resolution"],
            "" if config["which_dataset"] != "imagenet_lt" else "longtail",
            "_val" if config["split"] == "val" else "",
            "_test" if test_part else "",
            config["feature_extractor"],
            config["backbone_feature_extractor"],
        )
        print("Filenames are ", h5file_name_feats)

    print(
        "Starting to load dataset into an HDF5 file with chunk size %i and compression %s..."
        % (config["chunk_size"], config["compression"])
    )

    # Save original COCO image indexes in order for the evaluation set
    if config["which_dataset"] == "coco" and test_part:
        all_image_ids = []
    # Loop over loader
    for i, (x, y, image_id) in enumerate(tqdm(train_loader)):
        if config["which_dataset"] == "coco" and test_part:
            all_image_ids.append(image_id)
        if not config["save_images_only"]:
            with torch.no_grad():
                x_tf = x.cuda()
                x_tf = x_tf * 0.5 + 0.5
                x_tf = (x_tf - norm_mean) / norm_std
                x_tf = torch.nn.functional.upsample(x_tf, 224, mode="bicubic")

                x_feat, _ = net(x_tf)
                x_feat = x_feat.cpu().numpy()
                if config["feature_augmentation"]:
                    x_tf_hflip = tv_f.hflip(x_tf)
                    x_feat_hflip, _ = net(x_tf_hflip)
                    x_feat_hflip = x_feat_hflip.cpu().numpy()
                else:
                    x_feat_hflip = None
        else:
            x_feat, x_feat_hflip = None, None
        # Stick X into the range [0, 255] since it's coming from the train loader
        x = (255 * ((x + 1) / 2.0)).byte().numpy()
        # Numpyify y
        y = y.numpy()
        # If we're on the first batch, prepare the hdf5
        if i == 0:
            # Save images and labels in hdf5 file
            if not config["save_features_only"]:
                with h5.File(h5file_name, "w") as f:
                    print("Producing dataset of len %d" % len(train_loader.dataset))
                    imgs_dset = f.create_dataset(
                        "imgs",
                        x.shape,
                        dtype="uint8",
                        maxshape=(
                            len(train_loader.dataset),
                            3,
                            config["resolution"],
                            config["resolution"],
                        ),
                        chunks=(
                            config["chunk_size"],
                            3,
                            config["resolution"],
                            config["resolution"],
                        ),
                        compression=config["compression"],
                    )
                    print("Image chunks chosen as " + str(imgs_dset.chunks))
                    imgs_dset[...] = x
                    labels_dset = f.create_dataset(
                        "labels",
                        y.shape,
                        dtype="int64",
                        maxshape=(len(train_loader.dataset),),
                        chunks=(config["chunk_size"],),
                        compression=config["compression"],
                    )
                    print("Label chunks chosen as " + str(labels_dset.chunks))
                    labels_dset[...] = y

            # Save features in hdf5 file
            if not config["save_images_only"]:
                with h5.File(h5file_name_feats, "w") as f:
                    features_dset = f.create_dataset(
                        "feats",
                        x_feat.shape,
                        dtype="float",
                        maxshape=(len(train_loader.dataset), x_feat.shape[1]),
                        chunks=(config["chunk_size"], x_feat.shape[1]),
                        compression=config["compression"],
                    )
                    features_dset[...] = x_feat
                    if config["feature_augmentation"]:
                        features_dset_hflips = f.create_dataset(
                            "feats_hflip",
                            x_feat.shape,
                            dtype="float",
                            maxshape=(len(train_loader.dataset), x_feat.shape[1]),
                            chunks=(config["chunk_size"], x_feat.shape[1]),
                            compression=config["compression"],
                        )
                        features_dset_hflips[...] = x_feat_hflip

        # Else append to the hdf5
        else:
            if not config["save_features_only"]:
                with h5.File(h5file_name, "a") as f:
                    f["imgs"].resize(f["imgs"].shape[0] + x.shape[0], axis=0)
                    f["imgs"][-x.shape[0] :] = x
                    f["labels"].resize(f["labels"].shape[0] + y.shape[0], axis=0)
                    f["labels"][-y.shape[0] :] = y

            if not config["save_images_only"]:
                with h5.File(h5file_name_feats, "a") as f:
                    f["feats"].resize(f["feats"].shape[0] + x_feat.shape[0], axis=0)
                    f["feats"][-x_feat.shape[0] :] = x_feat
                if config["feature_augmentation"]:
                    with h5.File(h5file_name_feats, "a") as f:
                        f["feats_hflip"].resize(
                            f["feats_hflip"].shape[0] + x_feat_hflip.shape[0], axis=0
                        )
                        f["feats_hflip"][-x_feat_hflip.shape[0] :] = x_feat_hflip

    if config["which_dataset"] == "coco" and test_part:
        print(
            "Saved COCO index images for evaluation set (in order of appearance in the hdf5 file)"
        )
        np.save(
            os.path.join("coco_stuff_val_indexes", "cocostuff_val2_all_idxs"),
            np.concatenate(all_image_ids),
        )


def main():
    # parse command line and run
    parser = prepare_parser()
    config = vars(parser.parse_args())
    print(config)
    run(config)


if __name__ == "__main__":
    main()

	# Copyright (c) Facebook, Inc. and its affiliates.
	# All rights reserved.
	#
	# This source code is licensed under the license found in the
	# LICENSE file in the root directory of this source tree.
	#
	# All contributions by Andy Brock:
	# Copyright (c) 2019 Andy Brock
	#
	# MIT License
	""" Convert dataset to HDF5
	This script preprocesses a dataset and saves it (images and labels) to
	an HDF5 file for improved I/O. """
	import os
	import sys
	from argparse import ArgumentParser
	from tqdm import tqdm, trange
	import h5py as h5

	import numpy as np
	import torch

	import utils
	import torchvision.transforms.functional as tv_f


	def prepare_parser():
	usage = "Parser for ImageNet HDF5 scripts."
	parser = ArgumentParser(description=usage)
	parser.add_argument(
	"--resolution",
	type=int,
	default=128,
	help="Which Dataset resolution to train on, out of 64, 128, 256 (default: %(default)s)",
	)
	parser.add_argument(
	"--split",
	type=str,
	default="train",
	help="Which Dataset to convert: train, val (default: %(default)s)",
	)
	parser.add_argument(
	"--data_root",
	type=str,
	default="data",
	help="Default location where data is stored (default: %(default)s)",
	)
	parser.add_argument(
	"--out_path",
	type=str,
	default="data",
	help="Default location where data in hdf5 format will be stored (default: %(default)s)",
	)
	parser.add_argument(
	"--pretrained_model_path",
	type=str,
	default="",
	help="Location where the pretrained model (to extract features) can be found (default: %(default)s)",
	)
	parser.add_argument(
	"--save_features_only",
	action="store_true",
	default=False,
	help="Only save features in hdf5 file.",
	)
	parser.add_argument(
	"--save_images_only",
	action="store_true",
	default=False,
	help="Only save images and their labels in hdf5 file.",
	)
	parser.add_argument(
	"--feature_augmentation",
	action="store_true",
	default=False,
	help="Additioally store instance features with horizontally flipped input images.",
	)
	parser.add_argument(
	"--feature_extractor",
	type=str,
	default="classification",
	choices=["classification", "selfsupervised"],
	help="Choice of feature extractor",
	)
	parser.add_argument(
	"--backbone_feature_extractor",
	type=str,
	default="resnet50",
	choices=["resnet50"],
	help="Choice of feature extractor backbone",
	)
	parser.add_argument(
	"--which_dataset", type=str, default="imagenet", help="Dataset choice."
	)
	parser.add_argument(
	"--instance_json",
	type=str,
	default="",
	help="Path to JSON containing instance segmentations for COCO_Stuff",
	)
	parser.add_argument(
	"--stuff_json",
	type=str,
	default="",
	help="Path to JSON containing instance segmentations for COCO_Stuff",
	)
	parser.add_argument(
	"--batch_size",
	type=int,
	default=256,
	help="Default overall batchsize (default: %(default)s)",
	)
	parser.add_argument(
	"--num_workers",
	type=int,
	default=16,
	help="Number of dataloader workers (default: %(default)s)",
	)
	parser.add_argument(
	"--chunk_size",
	type=int,
	default=500,
	help="Default overall batchsize (default: %(default)s)",
	)
	parser.add_argument(
	"--compression",
	action="store_true",
	default=False,
	help="Use LZF compression? (default: %(default)s)",
	)
	return parser


	def run(config):
	# Load pretrained feature extractor
	net = utils.load_pretrained_feature_extractor(
	config["pretrained_model_path"],
	config["feature_extractor"],
	config["backbone_feature_extractor"],
	)
	net.eval()

	# Update compression entry
	config["compression"] = (
	"lzf" if config["compression"] else None
	) # No compression; can also use 'lzf'

	# Get dataset
	kwargs = {
	"num_workers": config["num_workers"],
	"pin_memory": False,
	"drop_last": False,
	}
	test_part = False
	if config["split"] == "test":
	config["split"] = "val"
	test_part = True
	if config["which_dataset"] in ["imagenet", "imagenet_lt"]:
	data_path = os.path.join(config["data_root"], config["split"])
	else:
	data_path = config["data_root"]
	dataset = utils.get_dataset_images(
	config["resolution"],
	data_path=data_path,
	longtail=config["which_dataset"] == "imagenet_lt",
	split=config["split"],
	test_part=test_part,
	which_dataset=config["which_dataset"],
	instance_json=config["instance_json"],
	stuff_json=config["stuff_json"],
	)
	train_loader = utils.get_dataloader(
	dataset, config["batch_size"], shuffle=False, **kwargs
	)

	# HDF5 supports chunking and compression. You may want to experiment
	# with different chunk sizes to see how it runs on your machines.
	# Chunk Size/compression Read speed @ 256x256 Read speed @ 128x128 Filesize @ 128x128 Time to write @128x128
	# 1 / None 20/s
	# 500 / None ramps up to 77/s 102/s 61GB 23min
	# 500 / LZF 8/s 56GB 23min
	# 1000 / None 78/s
	# 5000 / None 81/s
	# auto:(125,1,16,32) / None 11/s 61GB

	# Use imagenet statistics to preprocess images for the feature extractor (instance features)
	norm_mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda()
	norm_std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda()

	if config["which_dataset"] in ["imagenet", "imagenet_lt"]:
	dataset_name_prefix = "ILSVRC"
	elif config["which_dataset"] == "coco":
	dataset_name_prefix = "COCO"
	else:
	dataset_name_prefix = config["which_dataset"]

	if not config["save_features_only"]:
	h5file_name = config["out_path"] + "/%s%i%s%s%s_xy.hdf5" % (
	dataset_name_prefix,
	config["resolution"],
	"" if config["which_dataset"] != "imagenet_lt" else "longtail",
	"_val" if config["split"] == "val" else "",
	"_test" if test_part else "",
	)
	print("Filenames are ", h5file_name)

	if not config["save_images_only"]:
	h5file_name_feats = config["out_path"] + "/%s%i%s%s%s_feats_%s_%s.hdf5" % (
	dataset_name_prefix,
	config["resolution"],
	"" if config["which_dataset"] != "imagenet_lt" else "longtail",
	"_val" if config["split"] == "val" else "",
	"_test" if test_part else "",
	config["feature_extractor"],
	config["backbone_feature_extractor"],
	)
	print("Filenames are ", h5file_name_feats)

	print(
	"Starting to load dataset into an HDF5 file with chunk size %i and compression %s..."
	% (config["chunk_size"], config["compression"])
	)

	# Save original COCO image indexes in order for the evaluation set
	if config["which_dataset"] == "coco" and test_part:
	all_image_ids = []
	# Loop over loader
	for i, (x, y, image_id) in enumerate(tqdm(train_loader)):
	if config["which_dataset"] == "coco" and test_part:
	all_image_ids.append(image_id)
	if not config["save_images_only"]:
	with torch.no_grad():
	x_tf = x.cuda()
	x_tf = x_tf * 0.5 + 0.5
	x_tf = (x_tf - norm_mean) / norm_std
	x_tf = torch.nn.functional.upsample(x_tf, 224, mode="bicubic")

	x_feat, _ = net(x_tf)
	x_feat = x_feat.cpu().numpy()
	if config["feature_augmentation"]:
	x_tf_hflip = tv_f.hflip(x_tf)
	x_feat_hflip, _ = net(x_tf_hflip)
	x_feat_hflip = x_feat_hflip.cpu().numpy()
	else:
	x_feat_hflip = None
	else:
	x_feat, x_feat_hflip = None, None
	# Stick X into the range [0, 255] since it's coming from the train loader
	x = (255 * ((x + 1) / 2.0)).byte().numpy()
	# Numpyify y
	y = y.numpy()
	# If we're on the first batch, prepare the hdf5
	if i == 0:
	# Save images and labels in hdf5 file
	if not config["save_features_only"]:
	with h5.File(h5file_name, "w") as f:
	print("Producing dataset of len %d" % len(train_loader.dataset))
	imgs_dset = f.create_dataset(
	"imgs",
	x.shape,
	dtype="uint8",
	maxshape=(
	len(train_loader.dataset),
	3,
	config["resolution"],
	config["resolution"],
	),
	chunks=(
	config["chunk_size"],
	3,
	config["resolution"],
	config["resolution"],
	),
	compression=config["compression"],
	)
	print("Image chunks chosen as " + str(imgs_dset.chunks))
	imgs_dset[...] = x
	labels_dset = f.create_dataset(
	"labels",
	y.shape,
	dtype="int64",
	maxshape=(len(train_loader.dataset),),
	chunks=(config["chunk_size"],),
	compression=config["compression"],
	)
	print("Label chunks chosen as " + str(labels_dset.chunks))
	labels_dset[...] = y

	# Save features in hdf5 file
	if not config["save_images_only"]:
	with h5.File(h5file_name_feats, "w") as f:
	features_dset = f.create_dataset(
	"feats",
	x_feat.shape,
	dtype="float",
	maxshape=(len(train_loader.dataset), x_feat.shape[1]),
	chunks=(config["chunk_size"], x_feat.shape[1]),
	compression=config["compression"],
	)
	features_dset[...] = x_feat
	if config["feature_augmentation"]:
	features_dset_hflips = f.create_dataset(
	"feats_hflip",
	x_feat.shape,
	dtype="float",
	maxshape=(len(train_loader.dataset), x_feat.shape[1]),
	chunks=(config["chunk_size"], x_feat.shape[1]),
	compression=config["compression"],
	)
	features_dset_hflips[...] = x_feat_hflip

	# Else append to the hdf5
	else:
	if not config["save_features_only"]:
	with h5.File(h5file_name, "a") as f:
	f["imgs"].resize(f["imgs"].shape[0] + x.shape[0], axis=0)
	f["imgs"][-x.shape[0] :] = x
	f["labels"].resize(f["labels"].shape[0] + y.shape[0], axis=0)
	f["labels"][-y.shape[0] :] = y

	if not config["save_images_only"]:
	with h5.File(h5file_name_feats, "a") as f:
	f["feats"].resize(f["feats"].shape[0] + x_feat.shape[0], axis=0)
	f["feats"][-x_feat.shape[0] :] = x_feat
	if config["feature_augmentation"]:
	with h5.File(h5file_name_feats, "a") as f:
	f["feats_hflip"].resize(
	f["feats_hflip"].shape[0] + x_feat_hflip.shape[0], axis=0
	)
	f["feats_hflip"][-x_feat_hflip.shape[0] :] = x_feat_hflip

	if config["which_dataset"] == "coco" and test_part:
	print(
	"Saved COCO index images for evaluation set (in order of appearance in the hdf5 file)"
	)
	np.save(
	os.path.join("coco_stuff_val_indexes", "cocostuff_val2_all_idxs"),
	np.concatenate(all_image_ids),
	)


	def main():
	# parse command line and run
	parser = prepare_parser()
	config = vars(parser.parse_args())
	print(config)
	run(config)


	if __name__ == "__main__":
	main()