models

import numpy as np
import torch
import torch.nn as nn

from modules.base_mkcl import BaseMKCL
from modules.visual_extractor import VisualExtractor
from torch.nn.parameter import Parameter
from modules.contrastive_loss import SupConLoss

class Attention(nn.Module):

    def __init__(self, k_size, v_size, affine_size=512):
        super().__init__()
        self.affine_k = nn.Linear(k_size, affine_size, bias=False)
        self.affine_v = nn.Linear(v_size, affine_size, bias=False)
        self.affine = nn.Linear(affine_size, 1, bias=False)

    def forward(self, k, v):
        # k: batch size x hidden size
        # v: batch size x spatial size x hidden size
        # z: batch size x spatial size
        content_v = self.affine_k(k).unsqueeze(1) + self.affine_v(v)
        z = self.affine(torch.tanh(content_v)).squeeze(2)
        alpha = torch.softmax(z, dim=1)
        context = (v * alpha.unsqueeze(2)).sum(dim=1)
        return context, alpha


class BaseMKCLModel(nn.Module):
    def __init__(self, args, tokenizer, num_classes, fw_adj, bw_adj, feat_size=2048, embed_size=256, hidden_size=512): #, num_classes, fw_adj, bw_adj, feat_size=1024, embed_size=256, hidden_size=512
        super(BaseCMNModel, self).__init__()
        self.args = args
        self.num_classes = num_classes
        self.tokenizer = tokenizer
        self.fw_adj = fw_adj
        self.bw_adj = bw_adj
        self.feat_size = feat_size
        self.embed_size = embed_size
        self.hidden_size = hidden_size
        self.visual_extractor = VisualExtractor(args)
        self.encoder_decoder = BaseMKCL(args, tokenizer,)
  
    def __str__(self):
        model_parameters = filter(lambda p: p.requires_grad, self.parameters())
        params = sum([np.prod(p.size()) for p in model_parameters])
        return super().__str__() + '\nTrainable parameters: {}'.format(params)

    def forward_iu_xray(self, images, targets=None, mode='train', update_opts={}):
      
        att_feats_0, fc_feats_0 = self.visual_extractor(images[:, 0])
        att_feats_1, fc_feats_1 = self.visual_extractor(images[:, 1])
        fc_feats = torch.cat((fc_feats_0, fc_feats_1), dim=1)

      
        att_feats_0 = torch.cat((global_feats1.unsqueeze(1), att_feats_0), dim=1)
    
        att_feats_1 = torch.cat((global_feats2.unsqueeze(1), att_feats_1), dim=1)
    


        att_feats = torch.cat((att_feats_0, att_feats_1), dim=1)
        att_feats = self.zjxd(att_feats.transpose(1,2)).transpose(1,2)


        if mode == 'train':
            output = self.encoder_decoder(fc_feats, att_feats, targets, mode='forward')
            # c_loss = self.criterionContrastive(fc_feats_0, fc_feats_1)
            return output
        elif mode == 'sample':
            output, output_probs = self.encoder_decoder(fc_feats, att_feats, mode='sample', update_opts=update_opts)
            return output, output_probs
        else:
            raise ValueError