Spaces:
Sleeping
Sleeping
File size: 6,451 Bytes
b20c769 |
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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 |
import logging
import random
from copy import deepcopy
from typing import Dict, Optional
import torch
import torchvision.transforms.v2.functional as TVF
from einops import rearrange
from torch.utils.data import DataLoader
logger = logging.getLogger("__main__")
def get_embeddings(data_loader, model, device, subsample_tokens: Optional[float] = None):
embeddings = []
labels = []
if subsample_tokens:
print(f"Subsampling tokens with ratio {subsample_tokens}")
model = model.eval()
with torch.no_grad():
for batch in data_loader:
batch_labels = batch.pop("target")
if "s1" in batch:
batch["s1"] = batch["s1"].to(device).to(torch.bfloat16)
if "s2" in batch:
batch["s2"] = batch["s2"].to(device).to(torch.bfloat16)
if "months" in batch:
batch["months"] = batch["months"].to(device).long()
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
batch_embeddings = model(**batch) # (bsz, dim) or (bsz, tokens, dim)
if subsample_tokens is not None:
if len(batch_embeddings.shape) < 3:
raise ValueError("subsample tokens only works for segmentation tasks")
num_tokens_per_instance = batch_embeddings.shape[1]
num_instances_to_keep = int(num_tokens_per_instance * subsample_tokens)
sampled_indices = torch.randperm(num_tokens_per_instance)[:num_instances_to_keep]
batch_embeddings = batch_embeddings[:, sampled_indices]
tokens_per_dim = int(num_tokens_per_instance**0.5)
pixels_per_token_dim = int(batch_labels.shape[1] / tokens_per_dim)
batch_labels_per_token = rearrange(
batch_labels,
"b (t_h p_h) (t_w p_w) -> b (t_h t_w) (p_h p_w)",
t_h=tokens_per_dim,
t_w=tokens_per_dim,
p_h=pixels_per_token_dim,
p_w=pixels_per_token_dim,
)
batch_labels = batch_labels_per_token[:, sampled_indices]
embeddings.append(batch_embeddings.to(torch.bfloat16).cpu())
labels.append(batch_labels)
return torch.cat(embeddings, dim=0), torch.cat(labels, dim=0)
class DownstreamAugs(object):
"""
For now, lets have no parameters
Choose 1 of 8 transformations and apply it to space_x and the segmentation map (if needed)
"""
def __init__(self, enabled: bool):
self.enabled = enabled
self.transformations = [
self.no_transform, # No transformation
self.rotate_90, # 90-degree rotation
self.rotate_180, # 180-degree rotation
self.rotate_270, # 270-degree rotation
self.hflip, # Horizontal flip
self.vflip, # Vertical flip
self.hflip_rotate_90, # Horizontal flip of 90-degree rotated image
self.vflip_rotate_90, # Vertical flip of 90-degree rotated image
]
def no_transform(self, x):
return x
def rotate_90(self, x):
return TVF.rotate(x, 90)
def rotate_180(self, x):
return TVF.rotate(x, 180)
def rotate_270(self, x):
return TVF.rotate(x, 270)
def hflip(self, x):
return TVF.hflip(x)
def vflip(self, x):
return TVF.vflip(x)
def hflip_rotate_90(self, x):
return TVF.hflip(TVF.rotate(x, 90))
def vflip_rotate_90(self, x):
return TVF.vflip(TVF.rotate(x, 90))
def apply(self, image, target, task_type):
assert task_type in ["cls", "seg"]
# image is (H, W, C)
# target is either (1,) for classification or (H, W) for segmentation
if not self.enabled:
return image, target
# choose 1 of 8 possible augmentations
transformation = random.choice(self.transformations)
# transform image and rearrange
image = rearrange(image, "h w c -> c h w")
image = transformation(image)
image = rearrange(image, "c h w -> h w c")
if task_type == "cls":
return image, target
else:
# transform segmentation map and rearrange
assert target.shape[-1] == image.shape[-1]
assert target.shape[-2] == image.shape[-2]
target = rearrange(target, "h w -> 1 h w")
target = transformation(target)
target = rearrange(target, "1 h w -> h w")
return image, target
def get_loaders(
benchmark,
config,
model_name,
batch_size,
num_workers,
eval_type,
train_partition: Optional[str] = None,
valtest_partition: Optional[str] = None,
norm_ops: Optional[Dict] = None,
):
use_train_augs = True if eval_type == "FT" else False
dataclass_kwargs = deepcopy(benchmark["kwargs"])
if norm_ops is None:
dataclass_kwargs["norm_operation"] = config["models"][model_name]
else:
dataclass_kwargs["norm_operation"] = norm_ops
train_kwargs = deepcopy(dataclass_kwargs)
valtest_kwargs = deepcopy(dataclass_kwargs)
if train_partition is not None:
train_kwargs["partition"] = train_partition
if valtest_partition is None:
valtest_partition = "default"
valtest_kwargs["partition"] = valtest_partition
elif valtest_partition:
raise ValueError("Shouldn't have not None val_partition but None train_partiton")
return {
"train": DataLoader(
benchmark["class"](
**train_kwargs,
split="train",
augmentation=DownstreamAugs(use_train_augs),
),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
),
"valid": DataLoader(
benchmark["class"](
**valtest_kwargs,
split="valid",
augmentation=DownstreamAugs(False),
),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
),
"test": DataLoader(
benchmark["class"](
**valtest_kwargs,
split="test",
augmentation=DownstreamAugs(False),
),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
),
}
|