Spaces:
Sleeping
Sleeping
| #!/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"]) | |