Problem with extracting the feature

mathwseg · July 23, 2022, 2:14pm

I need to extract features for the images before classifying them… or removing the last layer for the classification model using vit-PyTorch

github.com

lucidrains/vit-pytorch/blob/main/vit_pytorch/max_vit.py

from functools import partial

import torch
from torch import nn, einsum

from einops import rearrange, repeat
from einops.layers.torch import Rearrange, Reduce

# helpers

def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

def cast_tuple(val, length = 1):
    return val if isinstance(val, tuple) else ((val,) * length)

# helper classes

This file has been truncated. show original

I tried to ignore the classification layer by
self.mlp_head = nn.Identity()

then do that code

from max_vit import MaxViT
from extractor import Extractor
model = MaxViT(
    num_classes = 0,
    dim = 192,                        
    depth = (2, 6, 14, 2)
    )          
feature= model(img)

but got the shape for each image torch.Size([1, 1536, 7, 7])
and got the file size 300MB for 1000 images ! … is there anything wrong with the code, please