TLM_Model.py

import torch
from torch import nn
from torch.nn import functional as F
from torch.distributions.normal import Normal
import math


import numpy as np
from efficientnet_pytorch.model import EfficientNet

N_CNN_OUT_CHANS =  1280 # 32 for My_CNN
# N_ACTIONS = 100 #10
# N_LOCS = 100 #16
N_CHANS = 3 # 3 for category, 9 for identity
# N_STEPS = 10 # 30
# EMBED_DIM = 256
DIM=4 # 4 for 128x128 output of the retina

# class My_CNN(nn.Module):
#     def __init__(self):
#         super(My_CNN, self).__init__()
#         self.conv1 = nn.Sequential(         
#             nn.Conv2d(
#                 in_channels=N_CHANS,      # 1 if mnist raw, no warpnet        
#                 out_channels=16,            
#                 kernel_size=5,              
#                 stride=1,                   
#                 padding=2,                  
#             ),                              
#             nn.ReLU(),                      
#             nn.MaxPool2d(kernel_size=2),    
#         )
#         self.conv2 = nn.Sequential(         
#             nn.Conv2d(16, 32, 5, 1, 2),     
#             nn.ReLU(),                      
#             nn.MaxPool2d(4),                
#         )
#         # fully connected layer, output 10 classes
#         self.out = nn.Linear(32 * DIM * DIM, N_ACTIONS )
        
       
        
#     def forward(self, x):
#         x = self.conv1(x)
#         x = self.conv2(x)
#         before_flatten_x = x
#         # flatten the output of conv2 to (batch_size, 32 * 8*8)
#         # x = x.view(x.size(0), -1)  
#         x = x.reshape(x.size(0), -1)  
        
#         return x, before_flatten_x #, self.out(x)  # return x for visualization
    

 
class GLIMPSE(nn.Module):

    ''''
    Glimpse network contains RETINA and an encoder.
    Encoder encodes output of RETINA and glimpse location.
    '''
    def __init__(self, embed_dim, device):
        super(GLIMPSE, self).__init__()

        self.fc_ro = nn.Linear(N_CNN_OUT_CHANS * DIM * DIM, embed_dim//2)  
        
        self.cnn =  EfficientNet.from_name(model_name='efficientnet-b0', in_channels=3) 

        loaing_path =  '/mnt/tempdata/Chen/epoch_285' #'/mnt/store1/motahareh/cnn_chckpt_new/chckpts_cnn_Mnih_RT_greyscale_imgnet' #'/mnt/store1/motahareh/all_model_chckpts_/model_chckpts_/chckpts_cnn_Mnih_RT_greyscale_imgnet_center_2'
        checkpoint = torch.load(loaing_path, weights_only=False)
        self.cnn.load_state_dict(checkpoint['model_state_dict'])
        with torch.no_grad():
            for name, param in self.cnn.named_parameters():

                if param.requires_grad:
                    param.requires_grad = False
                    if 'bn' in name:
                        # print("BATCHNORM name: ", name)
                        param.track_running_stats = False
            # for param in self.cnn.parameters():
            #     param.requires_grad = False
            #     print("param.name: ", param.names)
            #     if 'bn' in param.names:
            #         print("bn param.name: ", param.names)
            #         param.track_running_stats = False
        self.cnn = self.cnn.float().to(device)
         
            
        self.fc_lc = nn.Linear(2, embed_dim//2) # l -> hl
        
        self.fc_hg = nn.Linear(embed_dim//2,embed_dim) # f(hg)
        self.fc_hl = nn.Linear(embed_dim//2,embed_dim) # f(hl)
        

    def forward(self, x, l): # x: [batch_size, T, 3, g, g], g=50
        batch_size = x.size()[0]
        num_steps = x.size()[1]
        n_channels = x.size()[2]
        img_size = x.size()[3]
        # print("x.size(): ", x.size())
        cnn_in = x.view(batch_size * num_steps, n_channels, img_size, img_size)  # Flatten sequence dimension

        # print("cnn_in.size(): ", cnn_in.size())
        self.cnn.eval()
        final_feature_map = self.cnn.extract_features(cnn_in) # first one is logits that we discard
        # print("final_feature_map.size(): ", final_feature_map.size())
        before_flatten_x = final_feature_map
        # print("final_feature_map.size(): ", final_feature_map.size())
        ro = final_feature_map.flatten(start_dim=1)
        # print("ro.size(): ", ro.size())
        ro = ro.view(batch_size, num_steps, -1)
        # print("ro.size(): ", ro.size())
        # print("ro_seq.size(): ", ro.size())
        hg = F.relu(self.fc_ro(ro))                  # hg = fg(ro)
        # print("hg.size(): ", hg.size())
        # print("l.size(): ", l.size())
        hl = F.relu(self.fc_lc(l))                   # hl = fl(l)
        # print("hl.size(): ", hl.size())
        # hg_seq = hg.view(batch_size, num_steps, -1)
     
        # print("hg_seq.size(): ", hg_seq.size())
        g  = F.relu(self.fc_hg(hg)+self.fc_hl(hl))   # g = fg(hg,hl)
        # print("g.size(): ", g.size())
        return  g #, ro, hg #, class_labels # output: [batch_size, T, EMBED_DIM]
    
# Creates a square Sequential/Causal mask of size sz
def generate_causal_mask(sz: int):
    return torch.triu(torch.ones(sz, sz) * float('-inf'), diagonal=1)

class PositionalEncoding(nn.Module):
    # Positional encoding module taken from PyTorch Tutorial
    # Link: https://pytorch.org/tutorials/beginner/transformer_tutorial.html

    def __init__(self, d_model: int, max_len: int):
        super().__init__()
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        pe[:, 0, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """
        Args:
            x: Tensor, shape [seq_len, batch_size, embedding_dim]
        """
        x = x.permute(1,0,2)
        x = x + self.pe[:x.size(0)].to(x.device)
        x = x.permute(1,0,2)
        return x
    
class CausalTransformer(nn.Module):
    def __init__(self, hidden_size, nhead, n_blocks, n_steps):
        super().__init__()
        # Positional Encoder
        self.pos_emb = PositionalEncoding(hidden_size, n_steps)
        
        self.encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_size, nhead=nhead, batch_first=True)
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=n_blocks)
        self.causal_mask = generate_causal_mask(sz=n_steps)
        
    def forward(self, x):
        # padding_mask = self.generate_pad_mask(x)
        x = self.pos_emb(x)
        out = self.encoder(x)  # self.encoder(x, mask=self.causal_mask.to(x.device), src_key_padding_mask=padding_mask.to(x.device)) 
        return out


 # Generates a padding masks for each sequence in a batch
    def generate_pad_mask(self, batch):
        pad_tensor = torch.zeros((batch.shape[2])).to(batch.device) # padding with ones --> torch.ones
        mask = np.zeros((batch.shape[0],batch.shape[1]))
        for s in range(0, batch.shape[0]):
            for v in range(0, batch[s].shape[0]):
                new_s = torch.all(batch[s][v] == pad_tensor)
                mask[s][v] = new_s
        return torch.tensor(mask).bool().to(batch.device)





class CORE(nn.Module):
    '''
    Core network is a recurrent network which maintains a behavior state.
    '''
    def __init__(self, embed_dim):
        super(CORE, self).__init__()
        self.hidden_dim = embed_dim
        self.input_dim = embed_dim
        # forget gate components
        self.linear_forget_w1 = nn.Linear(self.input_dim, self.hidden_dim, bias=True)
        self.linear_forget_r1 = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
        self.sigmoid_forget = nn.Sigmoid()

        # input gate components
        self.linear_gate_w2 = nn.Linear(self.input_dim, self.hidden_dim, bias=True)
        self.linear_gate_r2 = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
        self.sigmoid_gate = nn.Sigmoid()

        # cell memory components
        self.linear_gate_w3 = nn.Linear(self.input_dim, self.hidden_dim, bias=True)
        self.linear_gate_r3 = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
        self.activation_gate = nn.Tanh()

        # out gate components
        self.linear_gate_w4 = nn.Linear(self.input_dim, self.hidden_dim, bias=True)
        self.linear_gate_r4 = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
        self.sigmoid_hidden_out = nn.Sigmoid()

        self.activation_final = nn.Tanh()
        
    def forget(self, x, h):
        x = self.linear_forget_w1(x)
        h = self.linear_forget_r1(h)
        return self.sigmoid_forget(x + h)

    def input_gate(self, x, h):

        # Equation 1. input gate
        x_temp = self.linear_gate_w2(x)
        h_temp = self.linear_gate_r2(h)
        i = self.sigmoid_gate(x_temp + h_temp)
        return i

    def cell_memory_gate(self, i, f, x, h, c_prev):
        x = self.linear_gate_w3(x)
        h = self.linear_gate_r3(h)
        # new information part that will be injected in the new context
        k = self.activation_gate(x + h)
        g = k * i

        # forget old context/cell info
        c = f * c_prev
        # learn new context/cell info
        c_next = g + c
        return c_next
    def out_gate(self, x, h):
        x = self.linear_gate_w4(x)
        h = self.linear_gate_r4(h)
        return self.sigmoid_hidden_out(x + h)
        
        
    def forward(self, x, h, c):
        n_samples = x.size()[0]
        # Equation 1. input gate
        i = self.input_gate(x, h)
        # Equation 2. forget gate
        f = self.forget(x, h)
        # Equation 3. updating the cell memory
        c_next = self.cell_memory_gate(i, f, x, h, c)
        # Equation 4. calculate the main output gate
        o = self.out_gate(x, h)
        # Equation 5. produce next hidden output
        h_next = o * self.activation_final(c_next)

        return h_next, c_next

class LOCATION(nn.Module):
    '''
    Location network learns policy for sensing locations.
    '''
    def __init__(self, embed_dim, n_locs):
        super(LOCATION, self).__init__()

        # self.std = std
        self.fc = nn.Linear(embed_dim, n_locs)

    def forward(self, h):
        l = self.fc(h)               # compute mean of Gaussian
        # l = torch.tanh(l)               # squeeze location to ensure sensing within the boundaries of an image
        return l # N_ACTIONS-dim location logits

class ACTION(nn.Module):
    '''
    Action network learn policy for task specific actions.
    In case of classification actions are possible classes.
    This network will be trained with supervised loss in case of classification.
    '''
    def __init__(self, embed_dim, n_actions):
        super(ACTION, self).__init__()

        self.fc = nn.Linear(embed_dim, n_actions)

    def forward(self, h):
        return self.fc(h)  # Do not apply softmax as loss function will take care of it

class TRANSFORMER_MODEL(nn.Module):
    '''
    Model combines all the previous elements
    '''
    def __init__(self, embed_dim,  n_heads, n_blocks, n_steps, n_grids_per_side, device):
        super(TRANSFORMER_MODEL, self).__init__()

        self.glimps = GLIMPSE(embed_dim, device)
        self.core   = CausalTransformer(embed_dim, n_heads, n_blocks, n_steps) # CORE()
        self.location = LOCATION(embed_dim, n_grids_per_side**2)
        self.n_steps = n_steps
  
    def forward(self, x, l): # x: [Batch_size, n_steps, IMG_SIZE, IMG_SIZE], l: [Batch_size, n_steps, 2]. Images are already warped based on l
        batch_size = x.size()[0]
        # print("x.size(): ", x.size())
        g = self.glimps(x, l.to(x.device)) # output: [batch_size, T, EMBED_DIM]
        
        transformer_out_tokens = self.core(g)       
        pred_target_locs = self.location(transformer_out_tokens[:,1:self.n_steps,:])  
        return pred_target_locs #, ro, hg # [batch_size, n_steps-1, N_ACTIONS]






# import torch
# from torch import nn
# from torch.nn import functional as F
# from torch.distributions.normal import Normal
# import math
# import numpy as np
# import os
# from efficientnet.model import EfficientNet

# N_CNN_OUT_CHANS =  1280 # 32 for My_CNN
# N_ACTIONS = 100 #10
# N_LOCS = 100 #16
# N_CHANS = 3 # 3 for category, 9 for identity
# N_STEPS = 10 # 30
# # EMBED_DIM = 256
# DIM=4

# # class My_CNN(nn.Module):
# #     def __init__(self):
# #         super(My_CNN, self).__init__()
# #         self.conv1 = nn.Sequential(         
# #             nn.Conv2d(
# #                 in_channels=N_CHANS,      # 1 if mnist raw, no warpnet        
# #                 out_channels=16,            
# #                 kernel_size=5,              
# #                 stride=1,                   
# #                 padding=2,                  
# #             ),                              
# #             nn.ReLU(),                      
# #             nn.MaxPool2d(kernel_size=2),    
# #         )
# #         self.conv2 = nn.Sequential(         
# #             nn.Conv2d(16, 32, 5, 1, 2),     
# #             nn.ReLU(),                      
# #             nn.MaxPool2d(4),                
# #         )
# #         # fully connected layer, output 10 classes
# #         self.out = nn.Linear(32 * DIM * DIM, N_ACTIONS )
        
       
        
# #     def forward(self, x):
# #         x = self.conv1(x)
# #         x = self.conv2(x)
# #         before_flatten_x = x
# #         # flatten the output of conv2 to (batch_size, 32 * 8*8)
# #         # x = x.view(x.size(0), -1)  
# #         x = x.reshape(x.size(0), -1)  
        
# #         return x, before_flatten_x #, self.out(x)  # return x for visualization
    
# def check_for_nans(tensor, name):
#     if torch.isnan(tensor).any():
#         print(f"NaNs found in {name}")
#     if torch.isinf(tensor).any():
#         print(f"Infinite values found in {name}")
 
# class GLIMPSE(nn.Module):

#     ''''
#     Glimpse network contains RETINA and an encoder.
#     Encoder encodes output of RETINA and glimpse location.
#     '''
#     def __init__(self, embed_dim, device):
#         super(GLIMPSE, self).__init__()

#         self.fc_ro = nn.Linear(N_CNN_OUT_CHANS*DIM*DIM, embed_dim//2)  
        
#         self.cnn =  EfficientNet.from_name(model_name='efficientnet-b0', in_channels=3) 

#         loaing_path =    # os.environ['SLURM_TMPDIR'] + '/chckpts_cnn_Mnih_RT_greyscale_imgnet_center_2' #'/home/motahareh/classification_project/model_chckpts_/chckpts_cnn_Mnih_RT_greyscale_imgnet_center_2'
#         checkpoint = torch.load(loaing_path, map_location='cpu')
#         self.cnn.load_state_dict(checkpoint['model_state_dict'])
#         with torch.no_grad():
#             for name, param in self.cnn.named_parameters():
#                 if param.requires_grad:
#                     param.requires_grad = False
#                     if 'bn' in name:
#                         print("BATCHNORM name: ", name)
#                         param.track_running_stats = False
#             # for param in self.cnn.parameters():
#             #     param.requires_grad = False
#             #     print("param.name: ", param.names)
#             #     if 'bn' in param.names:
#             #         print("bn param.name: ", param.names)
#             #         param.track_running_stats = False
#         self.cnn = self.cnn.to(device)
         
            
#         self.fc_lc = nn.Linear(2, embed_dim//2) # l -> hl
#         self.fc_hg = nn.Linear(embed_dim//2,embed_dim) # f(hg)
#         self.fc_hl = nn.Linear(embed_dim//2,embed_dim) # f(hl)
        

#     def forward(self, x, l): # x: [batch_size, T, 3, g, g], g=50
#         batch_size = x.size()[0]
#         # print("x.size(): ", x.size())
#         cnn_in =  torch.cat((x[:,0,:,:,:], 
#                              x[:,1,:,:,:], 
#                              x[:,2,:,:,:], 
#                              x[:,3,:,:,:],
#                              x[:,4,:,:,:], 
#                              x[:,5,:,:,:],
#                              x[:,6,:,:,:],
#                              x[:,7,:,:,:],
#                              x[:,8,:,:,:],
#                              x[:,9,:,:,:]), dim=0)
#                             #  x[:,10,:,:,:],
#                             #  x[:,11,:,:,:],
#                             #  x[:,12,:,:,:],
#                             #  x[:,13,:,:,:],
#                             #  x[:,14,:,:,:],
#                             #  x[:,15,:,:,:],
#                             #  x[:,16,:,:,:],
#                             #  x[:,17,:,:,:],
#                             #  x[:,18,:,:,:],
#                             #  x[:,19,:,:,:],
#                             #  x[:,20,:,:,:],
#                             #  x[:,21,:,:,:],
#                             #  x[:,22,:,:,:],
#                             #  x[:,23,:,:,:],
#                             #  x[:,24,:,:,:],
#                             #  x[:,25,:,:,:],
#                             #  x[:,26,:,:,:],
#                             #  x[:,27,:,:,:],
#                             #  x[:,28,:,:,:],
#                             #  x[:,29,:,:,:]), dim=0)
#         check_for_nans(cnn_in, "cnn_in")
#         self.cnn.eval()
#         _, final_feature_map, avg_pooled_feature_map = self.cnn(cnn_in) # first one is logits that we discard
#         check_for_nans(final_feature_map, "final_feature_map")
#         # print("final_feature_map.size(): ", final_feature_map.size())
#         before_flatten_x = final_feature_map
#         # print("final_feature_map.size(): ", final_feature_map.size())
#         ro = final_feature_map.flatten(start_dim=1, end_dim=3)
#         # print("ro.size(): ", ro.size())
#         hg = F.relu(self.fc_ro(ro))     
#         check_for_nans(hg, "in Glimpse: hg")
#         hl = F.relu(self.fc_lc(l))       
#         check_for_nans(hl, "in Glimpse: hl")
#         hg_seq = torch.cat((hg[0:batch_size, :].unsqueeze(1), 
#                             hg[batch_size:2*batch_size,:].unsqueeze(1), 
#                             hg[2*batch_size:3*batch_size,:].unsqueeze(1), 
#                             hg[3*batch_size:4*batch_size,:].unsqueeze(1), 
#                             hg[4*batch_size:5*batch_size,:].unsqueeze(1), 
#                             hg[5*batch_size:6*batch_size,:].unsqueeze(1), 
#                             hg[6*batch_size:7*batch_size,:].unsqueeze(1),
#                             hg[7*batch_size:8*batch_size,:].unsqueeze(1),
#                             hg[8*batch_size:9*batch_size,:].unsqueeze(1),
#                             hg[9*batch_size:10*batch_size,:].unsqueeze(1)), dim=1)
#                             # hg[10*batch_size:11*batch_size,:].unsqueeze(1),
#                             # hg[11*batch_size:12*batch_size,:].unsqueeze(1),
#                             # hg[12*batch_size:13*batch_size,:].unsqueeze(1),
#                             # hg[13*batch_size:14*batch_size,:].unsqueeze(1),
#                             # hg[14*batch_size:15*batch_size,:].unsqueeze(1),
#                             # hg[15*batch_size:16*batch_size,:].unsqueeze(1),
#                             # hg[16*batch_size:17*batch_size,:].unsqueeze(1),
#                             # hg[17*batch_size:18*batch_size,:].unsqueeze(1),
#                             # hg[18*batch_size:19*batch_size,:].unsqueeze(1),
#                             # hg[19*batch_size:20*batch_size,:].unsqueeze(1),
#                             # hg[20*batch_size:21*batch_size,:].unsqueeze(1),
#                             # hg[21*batch_size:22*batch_size,:].unsqueeze(1),
#                             # hg[22*batch_size:23*batch_size,:].unsqueeze(1),
#                             # hg[23*batch_size:24*batch_size,:].unsqueeze(1),
#                             # hg[24*batch_size:25*batch_size,:].unsqueeze(1),
#                             # hg[25*batch_size:26*batch_size,:].unsqueeze(1),
#                             # hg[26*batch_size:27*batch_size,:].unsqueeze(1),
#                             # hg[27*batch_size:28*batch_size,:].unsqueeze(1),
#                             # hg[28*batch_size:29*batch_size,:].unsqueeze(1),
#                             # hg[29*batch_size:30*batch_size,:].unsqueeze(1)), dim=1)
#         check_for_nans(hg_seq, "in Glimpse: hg_seq")
#         g  = F.relu(self.fc_hg(hg_seq)+self.fc_hl(hl))   # g = fg(hg,hl)
#         check_for_nans(g, "in Glimpse: g")
#         return  g, before_flatten_x #, class_labels # output: [batch_size, T, EMBED_DIM]
    
# # Creates a square Sequential/Causal mask of size sz
# def generate_causal_mask(sz: int):
#     return torch.triu(torch.ones(sz, sz) * float('-inf'), diagonal=1)

# class PositionalEncoding(nn.Module):
#     # Positional encoding module taken from PyTorch Tutorial
#     # Link: https://pytorch.org/tutorials/beginner/transformer_tutorial.html

#     def __init__(self, d_model: int, max_len: int = N_STEPS):
#         super().__init__()
#         position = torch.arange(max_len).unsqueeze(1)
#         div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
#         pe = torch.zeros(max_len, 1, d_model)
#         pe[:, 0, 0::2] = torch.sin(position * div_term)
#         pe[:, 0, 1::2] = torch.cos(position * div_term)
#         self.register_buffer('pe', pe)

#     def forward(self, x):
#         """
#         Args:
#             x: Tensor, shape [seq_len, batch_size, embedding_dim]
#         """
#         check_for_nans(x, "in PosEncoding, input: x")
#         x = x.permute(1,0,2)
#         check_for_nans(self.pe[:x.size(0)], "in PosEncoding, self.pe[:x.size(0)]")
#         x = x + self.pe[:x.size(0)].to(x.device)
#         check_for_nans(x, "in PosEncoding: after pos_embedding")
#         x = x.permute(1,0,2)
#         return x
    
# class CausalTransformer(nn.Module):
#     def __init__(self, hidden_size, nhead, n_blocks, n_steps):
#         super().__init__()
#         # Positional Encoder
#         self.pos_emb = PositionalEncoding(hidden_size, n_steps)
        
#         self.encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_size, nhead=nhead, batch_first=True)
#         self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=n_blocks)
#         # self.causal_mask = generate_causal_mask(sz=n_steps)
        
#     def forward(self, x):
#         # padding_mask = self.generate_pad_mask(x)
#         check_for_nans(x, "in CausalTransformer: input")
#         x = self.pos_emb(x)
#         check_for_nans(x, "in CausalTransformer: after pos_embedding")
#         out = self.encoder(x)  # self.encoder(x, mask=self.causal_mask.to(x.device), src_key_padding_mask=padding_mask.to(x.device)) 
#         check_for_nans(out, "in CausalTransformer: out")
#         return out


# #  # Generates a padding masks for each sequence in a batch
# #     def generate_pad_mask(self, batch):
# #         pad_tensor = torch.zeros((batch.shape[2])).to(batch.device) # padding with ones --> torch.ones
# #         mask = np.zeros((batch.shape[0],batch.shape[1]))
# #         for s in range(0, batch.shape[0]):
# #             for v in range(0, batch[s].shape[0]):
# #                 new_s = torch.all(batch[s][v] == pad_tensor)
# #                 mask[s][v] = new_s
# #         return torch.tensor(mask).bool().to(batch.device)





# class CORE(nn.Module):
#     '''
#     Core network is a recurrent network which maintains a behavior state.
#     '''
#     def __init__(self, embed_dim):
#         super(CORE, self).__init__()
#         self.hidden_dim = embed_dim
#         self.input_dim = embed_dim
#         # forget gate components
#         self.linear_forget_w1 = nn.Linear(self.input_dim, self.hidden_dim, bias=True)
#         self.linear_forget_r1 = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
#         self.sigmoid_forget = nn.Sigmoid()

#         # input gate components
#         self.linear_gate_w2 = nn.Linear(self.input_dim, self.hidden_dim, bias=True)
#         self.linear_gate_r2 = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
#         self.sigmoid_gate = nn.Sigmoid()

#         # cell memory components
#         self.linear_gate_w3 = nn.Linear(self.input_dim, self.hidden_dim, bias=True)
#         self.linear_gate_r3 = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
#         self.activation_gate = nn.Tanh()

#         # out gate components
#         self.linear_gate_w4 = nn.Linear(self.input_dim, self.hidden_dim, bias=True)
#         self.linear_gate_r4 = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
#         self.sigmoid_hidden_out = nn.Sigmoid()

#         self.activation_final = nn.Tanh()
        
#     def forget(self, x, h):
#         x = self.linear_forget_w1(x)
#         h = self.linear_forget_r1(h)
#         return self.sigmoid_forget(x + h)

#     def input_gate(self, x, h):

#         # Equation 1. input gate
#         x_temp = self.linear_gate_w2(x)
#         h_temp = self.linear_gate_r2(h)
#         i = self.sigmoid_gate(x_temp + h_temp)
#         return i

#     def cell_memory_gate(self, i, f, x, h, c_prev):
#         x = self.linear_gate_w3(x)
#         h = self.linear_gate_r3(h)
#         # new information part that will be injected in the new context
#         k = self.activation_gate(x + h)
#         g = k * i

#         # forget old context/cell info
#         c = f * c_prev
#         # learn new context/cell info
#         c_next = g + c
#         return c_next
#     def out_gate(self, x, h):
#         x = self.linear_gate_w4(x)
#         h = self.linear_gate_r4(h)
#         return self.sigmoid_hidden_out(x + h)
        
        
#     def forward(self, x, h, c):
#         n_samples = x.size()[0]
#         # Equation 1. input gate
#         i = self.input_gate(x, h)
#         # Equation 2. forget gate
#         f = self.forget(x, h)
#         # Equation 3. updating the cell memory
#         c_next = self.cell_memory_gate(i, f, x, h, c)
#         # Equation 4. calculate the main output gate
#         o = self.out_gate(x, h)
#         # Equation 5. produce next hidden output
#         h_next = o * self.activation_final(c_next)

#         return h_next, c_next

# class LOCATION(nn.Module):
#     '''
#     Location network learns policy for sensing locations.
#     '''
#     def __init__(self, embed_dim):
#         super(LOCATION, self).__init__()

#         # self.std = std
#         self.fc = nn.Linear(embed_dim, N_LOCS)

#     def forward(self, h):
#         check_for_nans(h, "in LOCATION: h")
#         l = self.fc(h)               # compute mean of Gaussian
#         # l = torch.tanh(l)               # squeeze location to ensure sensing within the boundaries of an image
#         check_for_nans(l, "in LOCATION: l")
#         return l # N_ACTIONS-dim location logits

# class ACTION(nn.Module):
#     '''
#     Action network learn policy for task specific actions.
#     In case of classification actions are possible classes.
#     This network will be trained with supervised loss in case of classification.
#     '''
#     def __init__(self, embed_dim):
#         super(ACTION, self).__init__()

#         self.fc = nn.Linear(embed_dim,N_ACTIONS)

#     def forward(self, h):
#         check_for_nans(h, "in ACTION: h")
#         output = self.fc(h)
#         check_for_nans(output, "in ACTION: output")
#         return output  # Do not apply softmax as loss function will take care of it

# class TRANSFORMER_MODEL(nn.Module):
#     '''
#     Model combines all the previous elements
#     '''
#     def __init__(self, embed_dim,  n_heads, n_blocks, n_steps, device):
#         super(TRANSFORMER_MODEL, self).__init__()

#         self.glimps = GLIMPSE(embed_dim, device)
#         self.core   = CausalTransformer(embed_dim, n_heads, n_blocks, n_steps) # CORE()
#         self.location = LOCATION(embed_dim)
#         self.n_steps = n_steps
  
#     def forward(self, x, l): # x: [Batch_size, n_steps, IMG_SIZE, IMG_SIZE], l: [Batch_size, n_steps, 2]. Images are already warped based on l
#         batch_size = x.size()[0]
#         check_for_nans(x, "in Transformer_Model: x")
#         check_for_nans(l, "in Transformer_Model: l")
#         g, feature_maps = self.glimps(x, l.to(x.device)) # output: [batch_size, T, EMBED_DIM]
#         check_for_nans(g, "in Transformer_Model: g")
#         check_for_nans(feature_maps, "in Transformer_Model: feature_maps")
#         transformer_out_tokens = self.core(g)       
#         check_for_nans(transformer_out_tokens, "in Transformer_Model: transformer_out_tokens")
#         pred_target_locs = self.location(transformer_out_tokens[:,1:self.n_steps,:])  
#         check_for_nans(pred_target_locs, "in Transformer_Model: pred_target_locs")
#         return pred_target_locs, feature_maps # [batch_size, n_steps-1, N_ACTIONS]