13 lines
455 B
Python
13 lines
455 B
Python
|
|
from model.backbone.PVT import PyramidVisionTransformer
|
|
import torch
|
|
from torch import nn
|
|
|
|
if __name__ == '__main__':
|
|
input=torch.randn(1,3,224,224)
|
|
model = PyramidVisionTransformer(
|
|
patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4], qkv_bias=True,
|
|
norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1])
|
|
output=model(input)
|
|
print(output.shape)
|