Missing features on "timm" package from pypi

[skr3178]

I downloaded package pip install timm
The code file for "_features.py" in the timm package is missing these lines of code:

Missing feature_take_indices which is used in vision_transformers.py

These lines are however included in the github repo. So the pypi package is an update behind i guess.

""" PyTorch Feature Extraction Helpers

A collection of classes, functions, modules to help extract features from models
and provide a common interface for describing them.

The return_layers, module re-writing idea inspired by torchvision IntermediateLayerGetter
https://github.com/pytorch/vision/blob/d88d8961ae51507d0cb680329d985b1488b1b76b/torchvision/models/_utils.py

Hacked together by / Copyright 2020 Ross Wightman
"""
from collections import OrderedDict, defaultdict
from copy import deepcopy
from functools import partial
from typing import Dict, List, Optional, Sequence, Set, Tuple, Union

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

from timm.layers import Format


__all__ = [
    'FeatureInfo', 'FeatureHooks', 'FeatureDictNet', 'FeatureListNet', 'FeatureHookNet', 'FeatureGetterNet',
    'feature_take_indices'
]


def _take_indices(
        num_blocks: int,
        n: Optional[Union[int, List[int], Tuple[int]]],
) -> Tuple[Set[int], int]:
    if isinstance(n, int):
        assert n >= 0
        take_indices = {x for x in range(num_blocks - n, num_blocks)}
    else:
        take_indices = {num_blocks + idx if idx < 0 else idx for idx in n}
    return take_indices, max(take_indices)


def _take_indices_jit(
        num_blocks: int,
        n: Union[int, List[int], Tuple[int]],
) -> Tuple[List[int], int]:
    if isinstance(n, int):
        assert n >= 0
        take_indices = [num_blocks - n + i for i in range(n)]
    elif isinstance(n, tuple):
        # splitting this up is silly, but needed for torchscript type resolution of n
        take_indices = [num_blocks + idx if idx < 0 else idx for idx in n]
    else:
        take_indices = [num_blocks + idx if idx < 0 else idx for idx in n]
    return take_indices, max(take_indices)


def feature_take_indices(
        num_blocks: int,
        indices: Optional[Union[int, List[int], Tuple[int]]] = None,
) -> Tuple[List[int], int]:
    if indices is None:
        indices = num_blocks  # all blocks if None
    if torch.jit.is_scripting():
        return _take_indices_jit(num_blocks, indices)
    else:
        # NOTE non-jit returns Set[int] instead of List[int] but torchscript can't handle that anno
        return _take_indices(num_blocks, indices)


def _out_indices_as_tuple(x: Union[int, Tuple[int, ...]]) -> Tuple[int, ...]:
    if isinstance(x, int):
        # if indices is an int, take last N features
        return tuple(range(-x, 0))
    return tuple(x)


OutIndicesT = Union[int, Tuple[int, ...]]

[Ross Wightman]

@skr3178 not sure how this is a bug, can't mix code on the main branch with a pip release... if you want functionality on main branch pip install git+https://github.com/rwightman/pytorch-image-models.git

I'm sill testing and changing a few more things before the next pip release...