Spaces:
Sleeping
Sleeping
File size: 3,011 Bytes
8d4ee22 |
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 |
import torch
from tqdm import tqdm
def get_metrics_for_model(train_loader, test_loader, model, metric_evaluator):
(features_train, feature_maps_train, outputs_train, features_test, feature_maps_test,
outputs_test, labels) = [], [], [], [], [], [], []
device = "cuda" if torch.cuda.is_available() else "cpu"
model.eval()
model = model.to(device)
training_transforms = train_loader.dataset.transform
train_loader.dataset.transform = test_loader.dataset.transform # Use test transform for train
train_loader = torch.utils.data.DataLoader(train_loader.dataset, batch_size=100, shuffle=False) # Turn off shuffling
print("Going in get metrics")
linear_matrix = model.linear.weight
entries = torch.nonzero(linear_matrix)
rel_features = torch.unique(entries[:, 1])
with torch.no_grad():
iterator = tqdm(enumerate(train_loader), total=len(train_loader))
for batch_idx, (data, target) in iterator:
xs1 = data.to("cuda")
output, feature_maps, final_features = model(xs1, with_feature_maps=True, with_final_features=True,)
outputs_train.append(output.to("cpu"))
features_train.append(final_features.to("cpu"))
labels.append(target.to("cpu"))
total = 0
correct = 0
iterator = tqdm(enumerate(test_loader), total=len(test_loader))
for batch_idx, (data, target) in iterator:
xs1 = data.to("cuda")
output, feature_maps, final_features = model(xs1, with_feature_maps=True,
with_final_features=True, )
feature_maps_test.append(feature_maps[:, rel_features].to("cpu"))
outputs_test.append(output.to("cpu"))
total += target.size(0)
_, predicted = output.max(1)
correct += predicted.eq(target.to("cuda")).sum().item()
print("test accuracy: ", correct / total)
features_train = torch.cat(features_train)
outputs_train = torch.cat(outputs_train)
feature_maps_test = torch.cat(feature_maps_test)
outputs_test = torch.cat(outputs_test)
labels = torch.cat(labels)
linear_matrix = linear_matrix[:, rel_features]
print("Shape of linear matrix: ", linear_matrix.shape)
all_metrics_dict = metric_evaluator(features_train, outputs_train,
feature_maps_test,
outputs_test, linear_matrix, labels)
result_dict = {"Accuracy": correct / total, "NFfeatures": linear_matrix.shape[1],
"PerClass": torch.nonzero(linear_matrix).shape[0] / linear_matrix.shape[0],
}
result_dict.update(all_metrics_dict)
print(result_dict)
# Reset Train transforms
train_loader.dataset.transform = training_transforms
return result_dict
|