Migrations
/
pytorch-image-models
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update davit.py

Browse Source
pull/1630/head
Fredo Guan 3 years ago
parent dbef70fc31
commit c579c62c32
1 changed files with 26 additions and 29 deletions
Show Stats Download Patch File Download Diff File

55
timm/models/davit.py
Unescape View File

@ -411,35 +411,32 @@ class DaViT(nn.Module):
for stage_id, stage_param in enumerate(self.architecture): for stage_id, stage_param in enumerate(self.architecture):
layer_offset_id = len(list(itertools.chain(*self.architecture[:stage_id]))) layer_offset_id = len(list(itertools.chain(*self.architecture[:stage_id])))
stage = nn.Sequential( stage = nn.Sequential([
nn.ModuleList([ nn.Sequential([
nn.Sequential( ChannelBlock(
nn.ModuleList([ dim=self.embed_dims[item],
ChannelBlock( num_heads=self.num_heads[item],
dim=self.embed_dims[item], mlp_ratio=mlp_ratio,
num_heads=self.num_heads[item], qkv_bias=qkv_bias,
mlp_ratio=mlp_ratio, drop_path=dpr[2 * (layer_id + layer_offset_id) + attention_id],
qkv_bias=qkv_bias, norm_layer=nn.LayerNorm,
drop_path=dpr[2 * (layer_id + layer_offset_id) + attention_id], ffn=ffn,
norm_layer=nn.LayerNorm, cpe_act=cpe_act
ffn=ffn, ) if attention_type == 'channel' else
cpe_act=cpe_act SpatialBlock(
) if attention_type == 'channel' else dim=self.embed_dims[item],
SpatialBlock( num_heads=self.num_heads[item],
dim=self.embed_dims[item], mlp_ratio=mlp_ratio,
num_heads=self.num_heads[item], qkv_bias=qkv_bias,
mlp_ratio=mlp_ratio, drop_path=dpr[2 * (layer_id + layer_offset_id) + attention_id],
qkv_bias=qkv_bias, norm_layer=nn.LayerNorm,
drop_path=dpr[2 * (layer_id + layer_offset_id) + attention_id], ffn=ffn,
norm_layer=nn.LayerNorm, cpe_act=cpe_act,
ffn=ffn, window_size=window_size,
cpe_act=cpe_act, ) if attention_type == 'spatial' else None
window_size=window_size, for attention_id, attention_type in enumerate(attention_types)]
) if attention_type == 'spatial' else None ) for layer_id, item in enumerate(stage_param)
for attention_id, attention_type in enumerate(attention_types)]) ])
) for layer_id, item in enumerate(stage_param)
])
)
self.stages.add_module(f'stage_{stage_id}', stage) self.stages.add_module(f'stage_{stage_id}', stage)

Write Preview
Loading…
Cancel
Save