Fine Tuning Heads

Not the talking heads

Linear Probing and Fine Tuning Heads

---

source

RNNProbingHead

def RNNProbingHead(
    c_in, input_size, hidden_size, n_classes, missing_channel_indices:NoneType=None, module:str='GRU',
    rnn_dropout:float=0.0, num_rnn_layers:int=1, pool:str='average', predict_every_n_patches:int=1,
    bidirectional:bool=True, affine:bool=False, pre_norm:bool=True, mlp_final_head:bool=False,
    linear_dropout:float=0.0
):

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

m = RNNProbingHead(c_in=7, 
                    pool='average', 
                    input_size = 384, 
                    missing_channel_indices=None,#,[0,1],
                    bidirectional=True,
                    affine=True, 
                    hidden_size=1200,
                    module='GRU',
                    n_classes=5,
                    predict_every_n_patches=5,
                    rnn_dropout=0.,
                    num_rnn_layers=1,
                    mlp_final_head=True,
                    pre_norm=True)
x = torch.randn((4*7,960,384))
m(x, return_softmax=True).shape

torch.Size([4, 5, 192])

batch_size = 2
n_vars = 7
max_len = 480
d_model = 384

# Create sequences of different lengths
seq_lens = torch.randint(50, max_len, (batch_size,))

# Create input tensors with different sequence lengths
x_list = [torch.randn(n_vars, length, d_model) for length in seq_lens]
x_nested = torch.nested.as_nested_tensor(x_list, layout=torch.jagged)
x_nested = flatten_dim_to_batch(x_nested, dim=1)  # flatten n_vars into batch dimension
y = m(x_nested, return_softmax=False)
y.shape

---

source

AttentiveClassifier

def AttentiveClassifier(
    embed_dim:int=768, num_heads:int=12, mlp_ratio:float=4.0, depth:int=1, norm_layer:type=LayerNorm,
    init_std:float=0.02, qkv_bias:bool=True, num_classes:int=1000, complete_block:bool=True, num_queries:int=1,
    affine:bool=False, c_in:int=7, per_channel:bool=False
):

Attentive Classifier

---

source

AttentivePooler

def AttentivePooler(
    num_queries:int=1, embed_dim:int=768, num_heads:int=12, mlp_ratio:float=4.0, depth:int=1,
    norm_layer:type=LayerNorm, init_std:float=0.02, qkv_bias:bool=True, complete_block:bool=True
):

Attentive Pooler

---

source

CrossAttentionBlock

def CrossAttentionBlock(
    dim, num_heads, mlp_ratio:float=4.0, qkv_bias:bool=False, act_layer:type=GELU, norm_layer:type=LayerNorm
):

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

---

source

CrossAttention

def CrossAttention(
    dim, num_heads:int=12, qkv_bias:bool=False
):

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

.. note:: As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

#  [bs x nvars x d_model x num_patch] 
x = torch.randn(4 * 7, 50, 256)
cls = AttentiveClassifier(embed_dim=256, c_in=7, num_queries=1, num_heads=2, mlp_ratio=4.0, depth=1, norm_layer=nn.LayerNorm, init_std=0.02, qkv_bias=True, num_classes=1, complete_block=False, per_channel=False)
cls(x).shape

torch.Size([28, 1, 256])

torch.Size([4, 1])

#  [bs x nvars x d_model x num_patch] 

cls_ = AttentiveClassifier(embed_dim=768, num_queries=1, num_heads=8, mlp_ratio=4.0, depth=1, norm_layer=nn.LayerNorm, init_std=0.02, qkv_bias=True, num_classes=5, complete_block=False)
sum(p.numel() for p in cls_.parameters())

2390021

#  [bs x nvars x d_model x num_patch] 
x = torch.randn(4, 7*10, 256)
cls_ = AttentiveClassifierNoMelt(embed_dim=256, num_queries=1, num_heads=2, mlp_ratio=4.0, depth=1, norm_layer=nn.LayerNorm, init_std=0.02, qkv_bias=True, num_classes=1, complete_block=True)
cls_(x).shape

torch.Size([4, 1])