clip_architectures.py

# -----------------------------------------------------------------------------
# MIT License
#
# Copyright (c) 2024 Ontolearn Team
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# -----------------------------------------------------------------------------

import torch, torch.nn as nn
import random
from typing import List
from ontolearn.nces_modules import *    

class LengthLearner_LSTM(nn.Module):
    """LSTM architecture"""
    def __init__(self, input_size, output_size, proj_dim=256, rnn_n_layers=2, drop_prob=0.2):
        super().__init__()
        self.name = 'LSTM'
        self.loss = nn.CrossEntropyLoss()
        self.lstm = nn.LSTM(input_size, proj_dim, rnn_n_layers, 
                            dropout=drop_prob, batch_first=True)
        self.dropout = nn.Dropout(drop_prob)
        self.fc1 = nn.Linear(2*proj_dim, proj_dim)
        self.fc2 = nn.Linear(proj_dim, proj_dim)
        self.fc3 = nn.Linear(proj_dim, output_size)  
    
    def forward(self, x1, x2):
        ''' Forward pass through the network.'''
        x1, _ = self.lstm(x1)
        x1 = x1.sum(1).contiguous().view(x1.shape[0], -1)
        x2, _ = self.lstm(x2)
        x2 = x2.sum(1).contiguous().view(x2.shape[0], -1)
        x = torch.cat([x1, x2], dim=-1)
        x = self.fc1(x)
        x = torch.selu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = x + torch.tanh(x)
        x = self.fc3(x)
        return x

class LengthLearner_GRU(nn.Module):
    """GRU architecture"""
    def __init__(self, input_size, output_size, proj_dim=256, rnn_n_layers=2, drop_prob=0.2):
        super().__init__()
        self.name = 'GRU'
        self.loss = nn.CrossEntropyLoss()
        self.gru = nn.GRU(input_size, proj_dim, rnn_n_layers, 
                            dropout=drop_prob, batch_first=True)
        self.dropout = nn.Dropout(drop_prob)
        self.fc1 = nn.Linear(2*proj_dim, proj_dim)
        self.fc2 = nn.Linear(proj_dim, proj_dim)
        self.fc3 = nn.Linear(proj_dim, output_size)

    def forward(self, x1, x2):
        ''' Forward pass through the network.'''
        x1, _ = self.gru(x1)
        x1 = x1.sum(1).contiguous().view(x1.shape[0], -1)
        x2, _ = self.gru(x2)
        x2 = x2.sum(1).contiguous().view(x2.shape[0], -1)
        x = torch.cat([x1, x2], dim=-1)
        x = self.fc1(x)
        x = torch.selu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        x = x + torch.tanh(x)
        x = self.fc3(x)
        return x
        

class LengthLearner_CNN(nn.Module):
    """CNN architecture"""
    def __init__(self, input_size, output_size, num_examples, proj_dim=256, kernel_size: list=[[5,7], [5,7]], stride: list=[[3,3], [3,3]], drop_prob=0.2):
        super().__init__()
        assert isinstance(kernel_size, list) and isinstance(kernel_size[0], list), "kernel size and stride must be lists of lists, e.g., [[5,7], [5,7]]"
        self.name = 'CNN'
        self.loss = nn.CrossEntropyLoss()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=4, kernel_size=(kernel_size[0][0], kernel_size[0][1]), stride=(stride[0][0], stride[0][1]), padding=(0,0))
        self.conv2 = nn.Conv2d(in_channels=4, out_channels=8, kernel_size=(kernel_size[1][0], kernel_size[1][1]), stride=(stride[1][0], stride[1][1]), padding=(0,0))
        self.dropout1d = nn.Dropout(drop_prob)
        self.dropout2d = nn.Dropout2d(drop_prob)
        conv_out_dim = 3536
        self.fc1 = nn.Linear(conv_out_dim, proj_dim)
        self.fc2 = nn.Linear(proj_dim, proj_dim)
        self.fc3 = nn.Linear(proj_dim, output_size)
        
    def forward(self, x1, x2):
        ''' Forward pass through the network.'''
        x1 = x1.unsqueeze(1)
        x2 = x2.unsqueeze(1)
        x = torch.cat([x1, x2], dim=-2)
        x = self.conv1(x)
        x = torch.selu(x)
        x = self.dropout2d(x)
        x = self.conv2(x)
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        x = torch.selu(x)
        x = self.dropout1d(x)
        x = self.fc2(x)
        x = x + torch.tanh(x)
        x = self.fc3(x)
        return x
    
        
class LengthLearner_SetTransformer(nn.Module):
    """SetTransformer architecture."""
    def __init__(self, input_size, output_size, proj_dim=256, num_heads=4, num_seeds=1, num_inds=32):
        super().__init__()
        self.name = 'SetTransformer'
        self.loss = nn.CrossEntropyLoss()
        self.enc = nn.Sequential(
                ISAB(input_size, proj_dim, num_heads, num_inds),
                ISAB(proj_dim, proj_dim, num_heads, num_inds))
        self.dec = nn.Sequential(
                PMA(proj_dim, num_heads, num_seeds),
                nn.Linear(proj_dim, output_size))

    def forward(self, x1, x2):
        ''' Forward pass through the network.'''
        x1 = self.enc(x1)
        x2 = self.enc(x2)
        x = torch.cat([x1, x2], dim=-2)
        x = self.dec(x).squeeze()
        return x