| import torch |
| from torch import nn, Tensor |
|
|
|
|
| class VecDyT(nn.Module): |
| def __init__(self, input_shape): |
|
|
| super().__init__() |
|
|
| self.alpha = nn.Parameter(torch.randn(input_shape)) |
|
|
| def forward(self, x): |
| x = torch.tanh(self.alpha * x) |
| return x |
|
|
|
|
| class VecDyGeluSine(nn.Module): |
| def __init__(self, input_shape): |
|
|
| super().__init__() |
|
|
| self.alpha = nn.Parameter(torch.randn(input_shape)) |
| self.beta = nn.Parameter(torch.randn(input_shape)) |
| self.gamma = nn.Parameter(torch.randn(1)) |
| self.etta = nn.Parameter(torch.randn(1)) |
| self.gelu = nn.GELU() |
|
|
| def forward(self, x): |
|
|
| x = self.gamma * self.gelu(self.alpha * x) + self.etta * torch.sin(self.beta * x) |
|
|
| return x |
|
|
| class FFUnit(nn.Module): |
| def __init__(self,dim): |
|
|
| super().__init__() |
|
|
| self.proj = nn.Linear(dim,dim,bias=False) |
| self.modulate = VecDyGeluSine(dim) |
|
|
|
|
| def forward(self, x): |
|
|
| u, v = x, x |
|
|
| u = self.modulate(u) |
| v = self.proj(v) |
| g = u * v |
|
|
| return g |
|
|
|
|
|
|
| class TTT(nn.Module): |
| def __init__(self, dim: int): |
|
|
| super(TTT, self).__init__() |
|
|
| self.mapping = nn.Linear(dim,dim,bias=False) |
|
|
| def forward(self, in_seq: Tensor) -> Tensor: |
|
|
|
|
| outs = [] |
|
|
| for seq in range(in_seq.size(1)): |
|
|
| state = in_seq[:,seq,:] |
| train_view = state + torch.randn_like(state) |
| label_view = state |
| loss = nn.functional.mse_loss(self.mapping(train_view), label_view) |
| grads = torch.autograd.grad( |
| loss, self.mapping.parameters(),create_graph=True) |
| with torch.no_grad(): |
| for param, grad in zip(self.mapping.parameters(), grads): |
|
|
| param -= 0.01 * grad |
|
|
| readout = self.mapping(in_seq[:,seq,:]).detach() |
| outs.append(readout) |
| out = torch.stack(outs, dim=1) |
|
|
| return out |
|
|
| class FFUnit_TTT(nn.Module): |
| def __init__(self,dim): |
|
|
| super().__init__() |
|
|
| self.proj = TTT(dim) |
| self.modulate = VecDyGeluSine(dim) |
|
|
|
|
| def forward(self, x): |
|
|
| u, v = x, x |
|
|
| u = self.modulate(u) |
| v = self.proj(v) |
| g = u * v |
|
|
| return g |
|
|
| class LiteTensorMapperBlock(nn.Module): |
| def __init__(self, dim, num_patch): |
|
|
| super().__init__() |
|
|
| self.norm_1 = VecDyT(dim) |
| self.norm_2 = VecDyT(dim) |
| self.memory = FFUnit_TTT(dim) |
| self.feedforward = FFUnit(dim) |
|
|
|
|
| def forward(self, x): |
|
|
|
|
| memorypath,residual = x, x |
|
|
| memorypath = self.norm_1(memorypath) |
|
|
| memorypath = self.memory(memorypath) |
|
|
| x = memorypath + residual |
|
|
| FFpath, residual = x, x |
|
|
| FFpath = self.norm_2(FFpath) |
|
|
| FFpath = self.feedforward(FFpath) |
|
|
| x = FFpath + residual |
|
|
| return x |
|
|
|
|
| class LiteTensorMapper(nn.Module): |
| def __init__(self, d_model,num_patch, num_layers): |
| super().__init__() |
|
|
| self.model = nn.Sequential( |
| *[LiteTensorMapperBlock(d_model,num_patch) for _ in range(num_layers)] |
| ) |
|
|
| def forward(self, x): |
|
|
| return self.model(x) |
|
|
|
|
