Tan_pytorch_segmentation/pytorch_segmentation/PV_Attention/A2-Attention.py

10 lines
229 B
Python

from attention.A2Atttention import DoubleAttention
import torch
from torch import nn
from torch.nn import functional as F
input=torch.randn(50,512,7,7)
a2 = DoubleAttention(512,128,128,True)
output=a2(input)
print(output.shape)