When I comment the line of transpose at line u = u.transpose(0,1) in class AttenModFullXL , training accuracy for a single batch reaches 1, but if it is uncommented, the model training accuracy stucks at 25 % and not improves. Why does this happens. Reshape cannot be done , because the purpose of transpose is different. Does transpose cause any problem in gradient flow?
import torch
import torch.nn as nn
from icecream import ic
import os
class Flatten(nn.Module):
def __init__(self, ):
super().__init__()
def forward(self, x):
bs = x.shape[0]
return x.reshape((bs, -1))
class SeparateLinear(torch.nn.Module):
def __init__(self, n, ins, out):
super().__init__()
self.Linears = {}
for i in range(n):
self.Linears[str(i)] = nn.Linear(ins, out)
self.Linears = nn.ModuleDict(self.Linears)
def forward(self, x):
outs = []
bs = x.shape[0]
x = x.reshape((16,bs,32,6,6))
for i in range(x.shape[0]):
a = self.Linears[str(i)](x[i])
outs.append(a)
ou = torch.stack(outs)
ou = ou.reshape(bs,16,32,6,16)
##ic(ou.shape)
return ou
class Self_attention(nn.Module):
def __init__(self,dp = 0.1):
super().__init__()
self.dropout = nn.Dropout(dp)
def forward(self,q,k,v,mask=None):
matmul_qk = q @ k.transpose(-1,-2)
scaled_attention_logits = matmul_qk/ torch.sqrt(torch.tensor(k.shape[-1],dtype=torch.float32))
attention_weights = torch.nn.functional.softmax(scaled_attention_logits,dim=-1)
attention_weights = self.dropout(attention_weights)
output = attention_weights @ v
return [output, attention_weights.detach().clone()]
class AttentionHead(nn.Module):
def __init__(self,d_model,d_features,dp ):
super().__init__()
self.attn = Self_attention(dp)
self.query_tfm = nn.Linear(d_model,d_features)
self.key_tfm = nn.Linear(d_model,d_features)
self.values_tfm = nn.Linear(d_model,d_features)
self.kmemory = None
self.vmemory = None
def forward(self,queries,key,values,mask=None,discard_mem = False ):
if discard_mem :
self.kmemory = None
self.vmemory = None
Q = self.query_tfm(queries)
K = self.key_tfm(key)
V = self.values_tfm(values)
dK = K.detach().clone()
dV = V.detach().clone()
if self.kmemory == None:
self.kmemory = dK
self.vmemory = dV
else:
K = torch.cat((K,self.kmemory),dim=1)
V = torch.cat((V,self.vmemory),dim=1)
self.kmemory = torch.cat((dK,self.kmemory),dim=1)# concating in sequence length
self.vmemory = torch.cat((dV,self.vmemory),dim=1)
# print("Memory appended",self.kmemory.shape)
x,att_weight = self.attn(Q,K,V,None)
return x,att_weight
def pop_last(self,n):
if self.kmemory.shape[1] == n*32:
self.kmemory = self.kmemory[:,:-32,:]
self.vmemory = self.vmemory[:,:-32,:]
def print_seq_length(self):
ic(self.kmemory.shape)
ic(self.vmemory.shape)
class MultiheadAttention(nn.Module):
def __init__(self,d_model,d_feature,n_heads,dp=0.1):
super().__init__()
self.d_model = d_model
self.d_feature = d_feature
self.n_heads = n_heads
assert d_model == d_feature * n_heads
self.attn_heads = nn.ModuleList([
AttentionHead(d_model, d_feature, dp) for _ in range(n_heads)
])
self.projection = nn.Linear(d_feature * n_heads, d_model)
def forward(self,queries,keys,values,discard_mem = False,mask=None):
comb = [attn(queries, keys, values, mask=mask,discard_mem = discard_mem) # (Batch, Seq, Feature)
for i, attn in enumerate(self.attn_heads)]
# log_size(x[0], "output of single head")
attentions = []
xs = []
for u,att in comb:
xs.append(u)
attentions.append(att)
# reconcatenate
x = torch.cat(xs, dim=-1) # (Batch, Seq, D_Feature * n_heads)
attentions = torch.cat(attentions,dim=-1)
# log_size(x, "concatenated output")
x = self.projection(x) # (Batch, Seq, D_Model)
# log_size(x, "projected output")
return x,attentions
def pop_last(self,n):
for i,attn in enumerate(self.attn_heads):
attn.pop_last(n)
class MultiheadAttentionXL(nn.Module):
def __init__(self,d_model,d_feature,n_heads,dp=0.1):
super().__init__()
self.d_model = d_model
self.d_feature = d_feature
self.n_heads = n_heads
assert d_model == d_feature * n_heads
self.attn_heads = nn.ModuleList([
AttentionHead(d_model, d_feature, dp) for _ in range(n_heads)
])
self.projection = nn.Linear(d_feature * n_heads, d_model)
def forward(self,queries,keys,values,discard_mem = False,mask=None):
comb = [attn(queries, keys, values, mask=mask,discard_mem = discard_mem) # (Batch, Seq, Feature)
for i, attn in enumerate(self.attn_heads)]
# log_size(x[0], "output of single head")
attentions = []
xs = []
for u,att in comb:
xs.append(u)
attentions.append(att)
# reconcatenate
x = torch.cat(xs, dim=-1) # (Batch, Seq, D_Feature * n_heads)
attentions = torch.cat(attentions,dim=-1)
# log_size(x, "concatenated output")
x = self.projection(x) # (Batch, Seq, D_Model)
# log_size(x, "projected output")
return x,attentions
def pop_last(self,n):
for i,attn in enumerate(self.attn_heads):
attn.pop_last(n)
class AttenModFullXL(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, 3, stride=1)
self.pool1 = nn.MaxPool2d(2, stride=2)
self.dropout1 = nn.Dropout(0.1)
self.conv2 = nn.Conv2d(32, 32, 3, stride=1)
self.pool2 = nn.MaxPool2d(2, stride=2)
self.dropout2 = nn.Dropout(0.1)
self.expand = nn.Linear(36,64)
self.relu1 = nn.ReLU()
self.dropout3 = nn.Dropout(0.1)
self.features = None
self.memory = []
dicts = {}
for i in range(16):
dicts["M-atten-" + str(i)] = MultiheadAttentionXL(64,16, 4, dp=0.2)
self.multiheads = nn.ModuleDict(dicts)
self.atten_dropout = nn.Dropout(0.2)
self.layernorm1 = nn.LayerNorm(64)
num_layers = 2
self.xl_heads = {}
self.norm_layers = {}
for i in range(num_layers):
self.xl_heads["XL-"+str(i)] = MultiheadAttentionXL(32*64,256,8,dp=0.2)
self.xl_heads = nn.ModuleDict(self.xl_heads)
for i in range(num_layers):
self.norm_layers["LN-"+str(i)] = nn.LayerNorm(32*64)
self.norm_layers = nn.ModuleDict(self.norm_layers)
self.flin1 = nn.Linear(16*2048,512)
self.flin2 = nn.Linear(512,128)
self.flin3 = nn.Linear(128,25)
self.att_mat = None
if torch.cuda.is_available():
self = self.cuda()
def pop_memory(self,n):
for i in range(len(self.multiheads)):
self.multiheads["M-atten-" + str(i)].pop_last(n)
for i in range(len(self.xl_heads)):
self.xl_heads["XL-"+str(i)].pop_last(n)
def forward(self,t,discard_mem = False):
t = t/255.0 # normaliztion added to the model forward itself
bs = t.shape[0]
if torch.cuda.is_available():
sliced = torch.zeros((t.shape[0],16, 32, 32)).cuda()
else:
sliced = torch.zeros((t.shape[0],16, 32, 32))
# ic(x.shape)
for i in range(4):
for j in range(4):
sliced[:,i * 4 + j] = t[:,i * 32:i * 32 + 32, j * 32:j * 32 + 32]
x = sliced
x = x.reshape((bs*16, 1, 32, 32))
u = self.conv1(x)
u = self.pool1(u)
self.features = u.detach().clone()
u = self.dropout1(u)
u = self.conv2(u)
u = self.pool2(u)
u = self.dropout2(u)
u = u.reshape((bs*16,32,36))
u = self.expand(u)
u = self.relu1(u)
# ic(u.shape)
u = u.reshape((16,bs,32,64))
attention_out =[]
self.att_mat = []
for i in range(len(self.multiheads)):
att_out, _att_mat = (self.multiheads["M-atten-" + str(i)](u[i] ,u[i], u[i],discard_mem))
self.att_mat.append(_att_mat.detach().clone())
attention_out.append(att_out)
with torch.no_grad():
self.att_mat = torch.stack(self.att_mat)
attention_out_stacked = torch.stack(attention_out)
u = u + self.atten_dropout(attention_out_stacked)
u = self.layernorm1(u)
u = u.reshape((16,bs,32*64))
# u = u.transpose(0,1)
for i in range(len(self.xl_heads.keys())):
att_out,_att_mat = self.xl_heads["XL-"+str(i)](u,u,u,discard_mem)
u = u+ torch.nn.functional.dropout(u)
u = self.norm_layers["LN-"+str(i)](u)
# att_out,_att_mat = self.xl_heads["XL-0"](u,u,u,discard_mem)
# u = u+ torch.nn.functional.dropout(u)
# u = self.norm_layers["LN-0"](u)
u = u.reshape((bs,16*2048))
u = self.flin1(u)
u = nn.functional.relu(u)
u = self.flin2(u)
u = nn.functional.relu(u)
u = self.flin3(u)
# ic(u.shape)
return u