File size: 2,872 Bytes
8d015d4 |
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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import torch
import torch.nn as nn
import timm
import numpy as np
class twins_svt_large(nn.Module):
def __init__(self, pretrained=True):
super().__init__()
self.svt = timm.create_model('twins_svt_large', pretrained=pretrained)
del self.svt.head
del self.svt.patch_embeds[2]
del self.svt.patch_embeds[2]
del self.svt.blocks[2]
del self.svt.blocks[2]
del self.svt.pos_block[2]
del self.svt.pos_block[2]
def forward(self, x, data=None, layer=2):
B = x.shape[0]
for i, (embed, drop, blocks, pos_blk) in enumerate(
zip(self.svt.patch_embeds, self.svt.pos_drops, self.svt.blocks, self.svt.pos_block)):
x, size = embed(x)
x = drop(x)
for j, blk in enumerate(blocks):
x = blk(x, size)
if j==0:
x = pos_blk(x, size)
if i < len(self.svt.depths) - 1:
x = x.reshape(B, *size, -1).permute(0, 3, 1, 2).contiguous()
if i == layer-1:
break
return x
def compute_params(self, layer=2):
num = 0
for i, (embed, drop, blocks, pos_blk) in enumerate(
zip(self.svt.patch_embeds, self.svt.pos_drops, self.svt.blocks, self.svt.pos_block)):
for param in embed.parameters():
num += np.prod(param.size())
for param in drop.parameters():
num += np.prod(param.size())
for param in blocks.parameters():
num += np.prod(param.size())
for param in pos_blk.parameters():
num += np.prod(param.size())
if i == layer-1:
break
for param in self.svt.head.parameters():
num += np.prod(param.size())
return num
class twins_svt_large_context(nn.Module):
def __init__(self, pretrained=True):
super().__init__()
self.svt = timm.create_model('twins_svt_large_context', pretrained=pretrained)
def forward(self, x, data=None, layer=2):
B = x.shape[0]
for i, (embed, drop, blocks, pos_blk) in enumerate(
zip(self.svt.patch_embeds, self.svt.pos_drops, self.svt.blocks, self.svt.pos_block)):
x, size = embed(x)
x = drop(x)
for j, blk in enumerate(blocks):
x = blk(x, size)
if j==0:
x = pos_blk(x, size)
if i < len(self.svt.depths) - 1:
x = x.reshape(B, *size, -1).permute(0, 3, 1, 2).contiguous()
if i == layer-1:
break
return x
if __name__ == "__main__":
m = twins_svt_large()
input = torch.randn(2, 3, 400, 800)
out = m.extract_feature(input)
print(out.shape)
|