Efficientnet (timm) extract features in FastAi V2

Deep Learning
1

Hi
how easy it is to get a features extract from a trained model.
Everything worked on the fastai 1st version and reset.
But there are problems

fastai 2.1.17
cuda 10.1
pytorch 1.7

train:

import timm
from wwf.vision.timm import *

model_names = timm.list_models('efficientnet*', pretrained=True)
pprint(model_names)

batch_tfms = [*aug_transforms(size=224, max_warp=0, flip_vert=True, max_lighting=0.2), Normalize.from_stats(*imagenet_stats)]

bs=256

src = (ImageDataLoaders.from_df(
        df=df, valid_col='is_valid', 
        path=path, 
        folder='train', 
        item_tfms=Resize(224),
        batch_tfms=batch_tfms, bs=bs,
    ))

learn = timm_learner(src, 'efficientnet_b2', 
                     loss_func=CrossEntropyLossFlat(),
                     metrics=[error_rate, accuracy], 
                     model_dir="/home/jupyter/")

learn.fit_one_cycle(6)
learn.export('/home/jupyter/efficientnet_b2_export.pkl')

predict:

learn = load_learner("/home/jupyter/efficientnet_b2_export.pkl", cpu=False)

dblock = DataBlock(blocks= (ImageBlock),
                   get_items=get_image_files,
                   item_tfms=Resize(224),
                   batch_tfms=batch_tfms,
                   )

path = Path("/home/jupyter/dataset/")
files = get_image_files(path)
dsets = dblock.datasets(path)
dls = dblock.dataloaders(path)

dl = dls.test_dl(files, bs=256, num_workers=4, with_labels=True, shuffle=False)

try this

class SaveFeatures():
    features=None
    def __init__(self, m): 
        self.hook = m.register_forward_hook(self.hook_fn)
        self.features = None
        print(self.hook)
    def hook_fn(self, module, input, output): 
        print(self)
        out = output.detach().cpu().numpy()
        print(out.shape)
        if isinstance(self.features, type(None)):
            self.features = out
            print(1)
        else:
            self.features = np.row_stack((self.features, out))
            print(2)
    def remove(self): 
        self.hook.remove()

sf = SaveFeatures(learn.model[1][4])

a1, target = learn.tta(dl=dl, n=1, beta=None, use_max=False)

array = np.array(sf.features)

produces a strange array of a different size.
I think there is a mixing of indices or something like that.

Now since EfficientNet models doesn’t support indexing?

or how to use it?
https://rwightman.github.io/pytorch-image-models/feature_extraction/

import torch
import timm
m = timm.create_model('xception41', pretrained=True)
o = m(torch.randn(2, 3, 299, 299))
print(f'Original shape: {o.shape}')
o = m.forward_features(torch.randn(2, 3, 299, 299))
print(f'Unpooled shape: {o.shape}')
1 Like