<jupyter_start><jupyter_text>Robustness Against Data Corruption<jupyter_code>import sys

# check whether run in Colab
root = "."
if "google.colab" in sys.modules:
    print("Running in Colab.")
    !pip3 install matplotlib
    !pip3 install einops==0.3.0
    !pip3 install timm==0.4.9
    !git clone https://github.com/xxxnell/how-do-vits-work.git
    root = "./how-do-vits-work"
    sys.path.append(root)
import os
import yaml
import copy
from pathlib import Path
import numpy as np

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
# from torch.utils.tensorboard import SummaryWriter

import models
# import ops.trains as trains
import ops.tests as tests
import ops.datasets as datasets
# import ops.schedulers as schedulers
# import ops.adversarial as adv
# config_path = "%s/configs/cifar10_vit.yaml" % root
config_path = "%s/configs/cifar100_vit.yaml" % root
# config_path = "%s/configs/imagenet_vit.yaml" % root

with open(config_path) as f:
    args = yaml.load(f)
    print(args)
dataset_args = copy.deepcopy(args).get("dataset")
train_args = copy.deepcopy(args).get("train")
val_args = copy.deepcopy(args).get("val")
model_args = copy.deepcopy(args).get("model")
optim_args = copy.deepcopy(args).get("optim")
env_args = copy.deepcopy(args).get("env")
dataset_train, dataset_test = datasets.get_dataset(**dataset_args, download=True)
dataset_name = dataset_args["name"]
num_classes = len(dataset_train.classes)

dataset_train = DataLoader(dataset_train, 
                           shuffle=True, 
                           num_workers=train_args.get("num_workers", 4), 
                           batch_size=train_args.get("batch_size", 128))
dataset_test = DataLoader(dataset_test, 
                          num_workers=val_args.get("num_workers", 4), 
                          batch_size=val_args.get("batch_size", 128))

print("Train: %s, Test: %s, Classes: %s" % (
    len(dataset_train.dataset), 
    len(dataset_test.dataset), 
    num_classes
))<jupyter_output><empty_output><jupyter_text>Load Pretrained Models Prepare the pretrained models (the cells below provide the snippets for ResNet-50 and ViT-Ti):<jupyter_code># download and load a pretrained model for CIFAR-100
url = "https://github.com/xxxnell/how-do-vits-work-storage/releases/download/v0.1/resnet_50_cifar100_691cc9a9e4.pth.tar"
path = "checkpoints/resnet_50_cifar100_691cc9a9e4.pth.tar"
models.download(url=url, path=path)

name = "resnet_50"
uid = "691cc9a9e4"
model = models.get_model(name, num_classes=num_classes,  # timm does not provide a ResNet for CIFAR
                         stem=model_args.get("stem", False))
map_location = "cuda" if torch.cuda.is_available() else "cpu"
checkpoint = torch.load(path, map_location=map_location)
model.load_state_dict(checkpoint["state_dict"])
import copy
import timm
import torch
import torch.nn as nn

# download and load a pretrained model for CIFAR-100
url = "https://github.com/xxxnell/how-do-vits-work-storage/releases/download/v0.1/vit_ti_cifar100_9857b21357.pth.tar"
path = "checkpoints/vit_ti_cifar100_9857b21357.pth.tar"
models.download(url=url, path=path)

model = timm.models.vision_transformer.VisionTransformer(
    num_classes=num_classes, img_size=32, patch_size=2,  # for CIFAR
    embed_dim=192, depth=12, num_heads=3, qkv_bias=False,  # for ViT-Ti 
)
model.name = "vit_ti"
uid = "9857b21357"
models.stats(model)
map_location = "cuda" if torch.cuda.is_available() else "cpu"
checkpoint = torch.load(path, map_location=map_location)
model.load_state_dict(checkpoint["state_dict"])<jupyter_output><empty_output><jupyter_text>Parallelize the given `moodel` by splitting the input:<jupyter_code>name = model.name
model = nn.DataParallel(model)
model.name = name<jupyter_output><empty_output><jupyter_text>Test model performance on in-domain data:<jupyter_code>gpu = torch.cuda.is_available()

model = model.cuda() if gpu else model.cpu()
metrics_list = []
for n_ff in [1]:
    print("N: %s, " % n_ff, end="")
    *metrics, cal_diag = tests.test(model, n_ff, dataset_test, verbose=False, gpu=gpu)
    metrics_list.append([n_ff, *metrics])<jupyter_output><empty_output><jupyter_text>Measure Predictive Performances on Corrupted Datasets We measure predictive performances on 75 corrupted datasets (= datasets corrupted by [15 different types](https://github.com/hendrycks/robustness/blob/master/assets/imagenet-c.png) with 5 levels of intensity each).<jupyter_code>n_ff = 1

gpu = torch.cuda.is_available()
val_args = copy.deepcopy(args).get("val")
dataset_args = copy.deepcopy(args).get("dataset")

model = model.cuda() if gpu else model.cpu()
metrics_c = { intensity: {} for intensity in range(1, 6) }
for intensity in range(1, 6):
    for ctype in datasets.get_corruptions():
        dataset_c = datasets.get_dataset_c(**dataset_args, ctype=ctype, intensity=intensity, download=True)
        dataset_c = DataLoader(dataset_c, 
                               num_workers=val_args.get("num_workers", 4), 
                               batch_size=val_args.get("batch_size", 128))
        print("Corruption type: %s, Intensity: %d, " % (ctype, intensity), end="")
        *metrics, cal_diag = tests.test(model, n_ff, dataset_c, verbose=False, gpu=gpu)
        metrics_c[intensity][ctype] = metrics

leaderboard_path = os.path.join("leaderboard", "logs", dataset_name, model.name)
Path(leaderboard_path).mkdir(parents=True, exist_ok=True)
metrics_dir = os.path.join(leaderboard_path, "%s_%s_%s_%s_corrupted.csv" % (dataset_name, model.name, uid, n_ff))
metrics_c_list = [[i, typ, *metrics] for i, typ_metrics in metrics_c.items() for typ, metrics in typ_metrics.items()]
tests.save_metrics(metrics_dir, metrics_c_list)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1231846019247594, 'fraction_numerical': 0.021172353455818024, 'mean_word_length': 4.642645607107601, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '596890', 'license': 'Apache-2.0', 'n_tokens_mistral': 2076, 'n_tokens_neox': 1959, 'n_words': 405, 'path': 'robustness.ipynb', 'repo_name': 'xxxnell/how-do-vits-work', 'size': '5715'}