wehub-resource-sync
164 lines
6.9 KiB
Python
164 lines
6.9 KiB
Python
from typing import Optional, Iterable, List, Union, Tuple
|
|
|
|
import torch
|
|
|
|
|
|
def verify_out_features_out_indices(
|
|
out_features: Optional[Iterable[str]],
|
|
out_indices: Optional[Iterable[int]],
|
|
stage_names: Optional[Iterable[str]],
|
|
):
|
|
"""
|
|
Verify that out_indices and out_features are valid for the given stage_names.
|
|
"""
|
|
if stage_names is None:
|
|
raise ValueError("Stage_names must be set for transformers backbones")
|
|
|
|
if out_features is not None:
|
|
if not isinstance(out_features, (list,)):
|
|
raise ValueError(f"out_features must be a list got {type(out_features)}")
|
|
if any(feat not in stage_names for feat in out_features):
|
|
raise ValueError(
|
|
f"out_features must be a subset of stage_names: {stage_names} got {out_features}"
|
|
)
|
|
if len(out_features) != len(set(out_features)):
|
|
raise ValueError(
|
|
f"out_features must not contain any duplicates, got {out_features}"
|
|
)
|
|
if out_features != (
|
|
sorted_feats := [feat for feat in stage_names if feat in out_features]
|
|
):
|
|
raise ValueError(
|
|
f"out_features must be in the same order as stage_names, expected {sorted_feats} got {out_features}"
|
|
)
|
|
|
|
if out_indices is not None:
|
|
if not isinstance(out_indices, (list, tuple)):
|
|
raise ValueError(
|
|
f"out_indices must be a list or tuple, got {type(out_indices)}"
|
|
)
|
|
# Convert negative indices to their positive equivalent: [-1,] -> [len(stage_names) - 1,]
|
|
positive_indices = tuple(
|
|
idx % len(stage_names) if idx < 0 else idx for idx in out_indices
|
|
)
|
|
if any(idx for idx in positive_indices if idx not in range(len(stage_names))):
|
|
raise ValueError(
|
|
f"out_indices must be valid indices for stage_names {stage_names}, got {out_indices}"
|
|
)
|
|
if len(positive_indices) != len(set(positive_indices)):
|
|
msg = f"out_indices must not contain any duplicates, got {out_indices}"
|
|
msg += (
|
|
f"(equivalent to {positive_indices}))"
|
|
if positive_indices != out_indices
|
|
else ""
|
|
)
|
|
raise ValueError(msg)
|
|
if positive_indices != tuple(sorted(positive_indices)):
|
|
sorted_negative = tuple(
|
|
idx
|
|
for _, idx in sorted(
|
|
zip(positive_indices, out_indices), key=lambda x: x[0]
|
|
)
|
|
)
|
|
raise ValueError(
|
|
f"out_indices must be in the same order as stage_names, expected {sorted_negative} got {out_indices}"
|
|
)
|
|
|
|
if out_features is not None and out_indices is not None:
|
|
if len(out_features) != len(out_indices):
|
|
raise ValueError(
|
|
"out_features and out_indices should have the same length if both are set"
|
|
)
|
|
if out_features != [stage_names[idx] for idx in out_indices]:
|
|
raise ValueError(
|
|
"out_features and out_indices should correspond to the same stages if both are set"
|
|
)
|
|
|
|
|
|
def _align_output_features_output_indices(
|
|
out_features: Optional[List[str]],
|
|
out_indices: Optional[Union[List[int], Tuple[int]]],
|
|
stage_names: List[str],
|
|
):
|
|
"""
|
|
Finds the corresponding `out_features` and `out_indices` for the given `stage_names`.
|
|
|
|
The logic is as follows:
|
|
- `out_features` not set, `out_indices` set: `out_features` is set to the `out_features` corresponding to the
|
|
`out_indices`.
|
|
- `out_indices` not set, `out_features` set: `out_indices` is set to the `out_indices` corresponding to the
|
|
`out_features`.
|
|
- `out_indices` and `out_features` not set: `out_indices` and `out_features` are set to the last stage.
|
|
- `out_indices` and `out_features` set: input `out_indices` and `out_features` are returned.
|
|
|
|
Args:
|
|
out_features (`List[str]`): The names of the features for the backbone to output.
|
|
out_indices (`List[int]` or `Tuple[int]`): The indices of the features for the backbone to output.
|
|
stage_names (`List[str]`): The names of the stages of the backbone.
|
|
"""
|
|
if out_indices is None and out_features is None:
|
|
out_indices = [len(stage_names) - 1]
|
|
out_features = [stage_names[-1]]
|
|
elif out_indices is None and out_features is not None:
|
|
out_indices = [stage_names.index(layer) for layer in out_features]
|
|
elif out_features is None and out_indices is not None:
|
|
out_features = [stage_names[idx] for idx in out_indices]
|
|
return out_features, out_indices
|
|
|
|
|
|
def get_aligned_output_features_output_indices(
|
|
out_features: Optional[List[str]],
|
|
out_indices: Optional[Union[List[int], Tuple[int]]],
|
|
stage_names: List[str],
|
|
) -> Tuple[List[str], List[int]]:
|
|
"""
|
|
Get the `out_features` and `out_indices` so that they are aligned.
|
|
|
|
The logic is as follows:
|
|
- `out_features` not set, `out_indices` set: `out_features` is set to the `out_features` corresponding to the
|
|
`out_indices`.
|
|
- `out_indices` not set, `out_features` set: `out_indices` is set to the `out_indices` corresponding to the
|
|
`out_features`.
|
|
- `out_indices` and `out_features` not set: `out_indices` and `out_features` are set to the last stage.
|
|
- `out_indices` and `out_features` set: they are verified to be aligned.
|
|
|
|
Args:
|
|
out_features (`List[str]`): The names of the features for the backbone to output.
|
|
out_indices (`List[int]` or `Tuple[int]`): The indices of the features for the backbone to output.
|
|
stage_names (`List[str]`): The names of the stages of the backbone.
|
|
"""
|
|
# First verify that the out_features and out_indices are valid
|
|
verify_out_features_out_indices(
|
|
out_features=out_features, out_indices=out_indices, stage_names=stage_names
|
|
)
|
|
output_features, output_indices = _align_output_features_output_indices(
|
|
out_features=out_features, out_indices=out_indices, stage_names=stage_names
|
|
)
|
|
# Verify that the aligned out_features and out_indices are valid
|
|
verify_out_features_out_indices(
|
|
out_features=output_features,
|
|
out_indices=output_indices,
|
|
stage_names=stage_names,
|
|
)
|
|
return output_features, output_indices
|
|
|
|
|
|
def find_pruneable_heads_and_indices(
|
|
heads: list[int],
|
|
n_heads: int,
|
|
head_size: int,
|
|
already_pruned_heads: set[int],
|
|
) -> tuple[set[int], torch.LongTensor]:
|
|
mask = torch.ones(n_heads, head_size, dtype=torch.bool)
|
|
|
|
heads = set(heads) - already_pruned_heads
|
|
|
|
for head in heads:
|
|
# Shift the head index left by however many smaller heads
|
|
# were already removed earlier.
|
|
shifted_head = head - sum(1 for h in already_pruned_heads if h < head)
|
|
mask[shifted_head] = False
|
|
|
|
index = torch.arange(n_heads * head_size)[mask.view(-1)].long()
|
|
return heads, index
|