Spaces:
Running
Running
File size: 2,022 Bytes
f74ae8e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
#!/usr/bin/python3
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import MultiheadAttention, GRU, Linear, LayerNorm, Dropout
class FFN(nn.Module):
def __init__(self, d_model, bidirectional=True, dropout=0):
super(FFN, self).__init__()
self.gru = GRU(d_model, d_model * 2, 1, bidirectional=bidirectional)
if bidirectional:
self.linear = Linear(d_model * 2 * 2, d_model)
else:
self.linear = Linear(d_model * 2, d_model)
self.dropout = Dropout(dropout)
def forward(self, x):
self.gru.flatten_parameters()
x, _ = self.gru(x)
x = F.leaky_relu(x)
x = self.dropout(x)
x = self.linear(x)
return x
class TransformerBlock(nn.Module):
def __init__(self, d_model, n_heads, bidirectional=True, dropout=0):
super(TransformerBlock, self).__init__()
self.norm1 = LayerNorm(d_model)
self.attention = MultiheadAttention(d_model, n_heads, dropout=dropout)
self.dropout1 = Dropout(dropout)
self.norm2 = LayerNorm(d_model)
self.ffn = FFN(d_model, bidirectional=bidirectional)
self.dropout2 = Dropout(dropout)
self.norm3 = LayerNorm(d_model)
def forward(self, x, attn_mask=None, key_padding_mask=None):
xt = self.norm1(x)
xt, _ = self.attention(xt, xt, xt,
attn_mask=attn_mask,
key_padding_mask=key_padding_mask)
x = x + self.dropout1(xt)
xt = self.norm2(x)
xt = self.ffn(xt)
x = x + self.dropout2(xt)
x = self.norm3(x)
return x
def main():
x = torch.randn(4, 64, 401, 201)
b, c, t, f = x.size()
x = x.permute(0, 3, 2, 1).contiguous().view(b, f * t, c)
transformer = TransformerBlock(d_model=64, n_heads=4)
x = transformer(x)
x = x.view(b, f, t, c).permute(0, 3, 2, 1)
print(x.size())
if __name__ == '__main__':
main()
|