File size: 3,866 Bytes
3589275 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | import numpy as np
import torch
import tqdm
from src import utils
from src.datasets.common import get_dataloader, maybe_dictionarize
from src.datasets.registry import get_dataset
from src.heads import get_classification_head
from src.linearize import LinearizedImageEncoder
from src.modeling import ImageClassifier
def eval_single_dataset(image_encoder, dataset_name, args):
classification_head = get_classification_head(args, dataset_name)
model = ImageClassifier(image_encoder, classification_head)
model.eval()
dataset = get_dataset(
dataset_name,
model.val_preprocess,
location=args.data_location,
batch_size=args.batch_size,
)
dataloader = get_dataloader(dataset, is_train=False, args=args, image_encoder=None)
device = args.device
with torch.no_grad():
top1, correct, n = 0.0, 0.0, 0.0
for _, data in enumerate(tqdm.tqdm(dataloader)):
data = maybe_dictionarize(data)
x = data["images"].to(device)
y = data["labels"].to(device)
logits = utils.get_logits(x, model)
pred = logits.argmax(dim=1, keepdim=True).to(device)
correct += pred.eq(y.view_as(pred)).sum().item()
n += y.size(0)
top1 = correct / n
metrics = {"top1": top1}
print(f"Done evaluating on {dataset_name}. Accuracy: {100*top1:.2f}%")
return metrics
def evaluate(image_encoder, args):
if args.eval_datasets is None:
return
per_dataset_results = {}
eval_datasets = (
args.eval_datasets
if args.control_dataset is None
else args.eval_datasets + [args.control_dataset]
)
for dataset_name in eval_datasets:
print("Evaluating on", dataset_name)
results = eval_single_dataset(image_encoder, dataset_name, args)
print(f"{dataset_name} Top-1 accuracy: {results['top1']:.4f}")
per_dataset_results[dataset_name + ":top1"] = results["top1"]
return per_dataset_results
def evaluate_task_vector_at_coef(
task_vector, pretrained_checkpoint, args, scaling_coef, posthoc_linearization=False
):
image_encoder = task_vector.apply_to(
pretrained_checkpoint, scaling_coef=scaling_coef
)
if posthoc_linearization:
pretrained_encoder = task_vector.apply_to(
pretrained_checkpoint, scaling_coef=0.0
)
image_encoder = LinearizedImageEncoder(
init_encoder=pretrained_encoder, image_encoder=image_encoder, args=args
)
coef_info = evaluate(image_encoder, args)
coef_info = add_normalized_accuracy(coef_info, args)
coef_info["avg_normalized_top1"] = np.mean(
[coef_info[dataset + ":normalized_top1"] for dataset in args.eval_datasets]
)
coef_info["avg_top1"] = np.mean(
[coef_info[dataset + ":top1"] for dataset in args.eval_datasets]
)
return coef_info
def evaluate_task_vector(
task_vector, pretrained_checkpoint, args, posthoc_linearization=False
):
info = {}
for scaling_coef in np.linspace(0.0, 1.0, args.n_eval_points):
print(f"Evaluating for scaling coefficient {scaling_coef:.2f}")
info[scaling_coef] = evaluate_task_vector_at_coef(
task_vector,
pretrained_checkpoint,
args,
scaling_coef,
posthoc_linearization,
)
return info
def add_normalized_accuracy(results, args):
for dataset_name in args.eval_datasets:
results[dataset_name + ":normalized_top1"] = (
results[dataset_name + ":top1"] / args.finetuning_accuracies[dataset_name]
)
return results
def nonlinear_advantage(acc_linear, acc_nonlinear, num_classes):
err_linear = 1 - acc_linear
err_nonlinear = 1 - acc_nonlinear
return (err_linear - err_nonlinear) * num_classes / (num_classes - 1)
|