Tan_pytorch_segmentation/pytorch_segmentation/PV_Attention/Shuffle-Attention-Usage.py

11 lines
243 B
Python
Raw Permalink Normal View History

2025-05-19 20:48:24 +08:00
from attention.ShuffleAttention import ShuffleAttention
import torch
from torch import nn
from torch.nn import functional as F
input = torch.randn(50, 512, 7, 7)
se = ShuffleAttention(channel=512, G=8)
output = se(input)
print(output.shape)