-
Notifications
You must be signed in to change notification settings - Fork 2
/
data_prune.py
116 lines (101 loc) · 3.06 KB
/
data_prune.py
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
from dataclasses import dataclass
from typing import Union
import numpy as np
import torch
from bns import (
Prune,
PruneMode,
PruneType,
)
from tqdm import tqdm
from keras.applications import ResNet152V2
from keras.applications.resnet_v2 import preprocess_input
from distance import (
_get_lpips_model,
normalize_flatten_features,
)
@dataclass
class ExtractedFeatures:
features: Union[np.ndarray, list]
image_names: list[str]
def get_extracted_image_features(
dataset: torch.Tensor,
) -> ExtractedFeatures:
extracted_features, image_name = [], []
feature_extractor = ResNet152V2(
weights="imagenet",
include_top=False,
)
for image_index in tqdm(dataset):
img = np.expand_dims(image_index, axis=0)
img = preprocess_input(img)
features = feature_extractor.predict(img)
features = features.flatten()
extracted_features.append(features)
image_name.append(image_index)
return ExtractedFeatures(
features=extracted_features,
image_names=image_name,
)
def get_extracted_image_features_perceptual_(
dataset: torch.Tensor,
) -> ExtractedFeatures:
extracted_features, image_name = [], []
model = _get_lpips_model()
model = model.to(torch.device("cpu"))
dataset = dataset.permute(0, 3, 2, 1)
for image_index in tqdm(dataset):
l1, l2, l3, l4, l5 = model.features(image_index)
features = model.classifier(l5)
extracted_features.append(features)
image_name.append(image_index)
return ExtractedFeatures(
features=extracted_features,
image_names=image_name,
)
def get_extracted_image_features_perceptual1(
dataset: torch.Tensor,
) -> ExtractedFeatures:
extracted_features, image_name = [], []
model = _get_lpips_model()
model = model.to(torch.device("cpu"))
dataset = dataset.permute(0, 3, 2, 1)
for image_index in tqdm(dataset):
features = model.features(image_index)
activations_x = normalize_flatten_features(features)
extracted_features.append(activations_x)
image_name.append(image_index)
return ExtractedFeatures(
features=extracted_features,
image_names=image_name,
)
def get_prune_data(
data_name: str,
dataset,
labels,
ssl_type: str,
prune_fraction: float,
prune_mode: str = PruneMode.EASY, # pegged
prune_type: str = PruneType.BOTH, # pegged
random_state: int = 40,
feature_extraction: bool = False,
):
match (data_name, feature_extraction):
case ("cifar10", True):
data_features = get_extracted_image_features_perceptual1(
dataset=dataset,
).features
case _:
data_features = dataset
prune: Prune = Prune(
prune_type=True,
ssl_type=ssl_type,
number_cluster=len(np.unique(labels)),
random_state=random_state,
dataset=data_features,
)
return prune.get_prune(
prune_fraction=prune_fraction,
pruned_mode=prune_mode,
prune_type=prune_type,
)