Package: dataset-interfaces
Dataset InterfacesThis repository contains the code for our recent work:Dataset Interfaces: Diagnosing Model Failures Using Controllable Counterfactual GenerationJoshua Vendrow*, Saachi Jain*, Logan Engstrom, Aleksander MadryPaper:https://arxiv.org/abs/2302.07865Blog post:https://gradientscience.org/dataset-interfaces/Getting startedInstall using pip, or clone our repository.pip install dataset-interfacesExample:For a walkthrough of codebase, check out ourexample notebook. This notebook shows how to
construct a dataset interface for a subset of ImageNet and generate counterfactual examples.Before runningrun_textual_inverson, initialize an🤗Accelerateenvironment with:accelerateconfigConstructing a Dataset InterfaceConstructing a dataset interface consists or learning aclass tokenfor each class in a datset, which can then be included in textual prompts.To learn a single token, we use the following function:fromdataset_interfacesimportrun_textual_inversionembed=run_textual_inversion(train_path=train_path,# path to directory with training set for a single classtoken=token,# text to use for new token, e.g "<plate>"class_name=class_name,# natrual language class description, e.g., "plate")Once all the class tokens are learned, we can create a custom tokenizer and text encoder with these tokens:importinference_utilsasinfer_utilsinfer_utils.create_encoder(embeds=embeds,# list of learned embeddings (from the code block above)tokens=tokens,# list of token stringsclass_names=class_names,# list of natural language class descriptionsencoder_root=encoder_root# path where to store the tokenizer and encoder)Generating Counterfactual ExamplesWe can now generate counterfactual examples by incorporating our learned tokens in textual prompts. Thegeneratefunction generates images for a specific class in the dataset (indexed in the order that classes are passed when constructing the encoder). When specifying the text prompt, "" acts as a placeholder for the class token.fromdataset_interfacesimportgenerategenerate(encoder_root=encoder_root,c=c,# index of a specific classprompts="a photo of a <TOKEN> in the grass",# can be a single prompt or a list of promptsnum_samples=10,random_seed=0# no seed by default)CLIP MetricTo directly evaluate the quality of the generated image, we useCLIP similarityto quantify the presence of the object of interest and desired distribution shift in the image.We can measure CLIP similarity between a set of generated images and a given caption as follows:sim_class = infer_utils.clip_similarity(imgs, "a photo of a dog")
sim_shift = infer_utils.clip_similarity(imgs, "a photo in the grass")ImageNet* BenchmarkOur benchmark for the ImageNet dataset consists of two components: our 1,000 learned class tokens for ImageNet, and the images generated by these tokens in 23 distribution shifts.ImageNet* TokensThe 1,000 learned tokens are avaiable onHuggingFaceand can be downloaded with:wget https://huggingface.co/datasets/madrylab/imagenet-star-tokens/resolve/main/tokens.zipTo generate images with these tokens, we first create a text encoder with the tokens, which we use to seamlessly integrate the tokens in text prompts:token_path="./tokens".# path to the tokens from HuggingFaceinfer_utils.create_imagenet_star_encoder(token_path,encoder_root="./encoder_root_imagenet")Now, we can generate counterfactual examples of ImageNet from a textual prompt (See theexample notebookfor a walk-through):fromdataset_interfacesimportgenerateencoder_root="./encoder_root_imagenet"c=207# the class for golden retrieverprompt="a photo of a <TOKEN> wearing a hat"generate(encoder_root,c,prompt,num_samples=10)ImageNet* ImagesOur benchmark contains images in 23 distribution shifts, with 50k images per shift (50 per class for 1000 classes). These images are also available onHuggingFace. In this repo we also provide masks for each distribution shift indicating which images we filter out with our CLIP metrics, atmasks.npy.We provide a wrapper on toptorchvision.datasets.ImageFolderto construct a dataset object that filters the images o=un the benchmark using this mask. So, we can make a dataset object for a shift as follows:fromdataset_interfacesimportutilsroot="./imagenet-star"# the path where the dataset from HuggingFacemask_path="./masks.npy"# the path to the mask fileshift="in_the_snow"# the distribution shift of interestds=utils.ImageNet_Star_Dataset(root,shift=shift,mask_path=mask_path)CitationTo cite this paper, please use the following BibTex entry:@inproceedings{vendrow2023dataset,
   title = {Dataset Interfaces: Diagnosing Model Failures Using Controllable Counterfactual Generation},
   author = {Joshua Vendrow and Saachi Jain and Logan Engstrom and Aleksander Madry}, 
   booktitle = {ArXiv preprint arXiv:2302.07865},
   year = {2023}
}Maintainers:Josh VendrowSaachi Jain