Spaces:
Sleeping
Sleeping
File size: 2,761 Bytes
1d7b87c | 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 | #!/usr/bin/env python3
# -*- coding: utf-8 -*-
import torch.nn as nn
import torch
from tqdm import tqdm
import numpy as np
from functools import partial
from timm.models.vision_transformer import Block
import os
import timm
timm.layers.set_fused_attn(False)
class BaseModel(nn.Module):
def get_head_weights(self):
params = torch.load(os.path.join(os.path.dirname(__file__), "models/crop_classifier_half.pt"), map_location="cpu", weights_only=True)
params = {k: params[k] for k in ['norm.weight', 'norm.bias', 'head.weight', 'head.bias']}
self.params = params
class IndepClassifier(BaseModel):
def __init__(self, embed_dim=1024):
super().__init__()
self.get_head_weights()
self.head = nn.Linear(self.params["head.weight"].shape[1], self.params["head.weight"].shape[0])
self.load_state_dict(self.params, strict=False)
def forward(self, x):
x = self.head(x)
return x
class MCAClassifier(BaseModel):
def __init__(self, depth=2, num_heads=8):
super().__init__()
self.get_head_weights()
self.embed_dim = self.params["head.weight"].shape[1]
self.blocks = nn.Sequential(*[Block(dim=self.embed_dim, num_heads=num_heads,
norm_layer=partial(nn.LayerNorm, eps=1e-6))for i in range(depth)])
self.norm = nn.LayerNorm(self.embed_dim, eps=1e-6)
self.num_classes = self.params["head.weight"].shape[0]
self.head = nn.Linear(self.embed_dim, self.num_classes)
self.load_state_dict(self.params, strict=False)
def forward(self, x):
x = self.forward_blocks(x)
x = self.forward_head(x)
return x
def forward_blocks(self, x):
x = self.blocks(x)
return x
def forward_head(self, x):
x = self.head(x)
return x
def freeze_head(self):
self.norm.weight.requires_grad = False
self.norm.bias.requires_grad = False
self.head.weight.requires_grad = False
self.head.bias.requires_grad = False
if __name__ == "__main__":
from utils import crop_to_seq, load_data, filter_df, class_list
from dataset import Dataset
device = "cpu"
model = MCAClassifier(depth=1)
model.to(device)
df_imgs, scores, embeddings = load_data(os.path.join(os.path.dirname(__file__), "../data/serengeti_toy"))
df_imgs, df = filter_df(df_imgs, classes=class_list)
df = crop_to_seq(df_imgs)
dataset = Dataset(df, embeddings, class_list, min_seqlength=12, fixed_length=True, nb_synthetic_seq=[1, 2, 3])
dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, num_workers=1)
output_dict = model.predict(dataloader, device)
print(output_dict["acc"])
|