(self, in_chans=3, out_dim=192,
embed_dims=[96, 192, 384, 768], depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
init_values=[1, 1, 1, 1], heads_ranges=[3, 3, 3, 3], mlp_ratios=[3, 3, 3, 3], split_sizes=[1, 2, 2, 4],
sr_ratios=[8, 4, 2, 1], drop_path_rate=0.1, norm_layer=nn.LayerNorm,
patch_norm=True, use_checkpoints=[False, False, False, False],
mixer_types=['Local1', 'Local1', 'Global2', 'Global2'],
chunkwise_recurrents=[True, True, False, False],
layerscales=[False, False, False, False], layer_init_values=1e-6)
| 880 | |
| 881 | class VTPTRNet(nn.Module): |
| 882 | def __init__(self, in_chans=3, out_dim=192, |
| 883 | embed_dims=[96, 192, 384, 768], depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], |
| 884 | init_values=[1, 1, 1, 1], heads_ranges=[3, 3, 3, 3], mlp_ratios=[3, 3, 3, 3], split_sizes=[1, 2, 2, 4], |
| 885 | sr_ratios=[8, 4, 2, 1], drop_path_rate=0.1, norm_layer=nn.LayerNorm, |
| 886 | patch_norm=True, use_checkpoints=[False, False, False, False], |
| 887 | mixer_types=['Local1', 'Local1', 'Global2', 'Global2'], |
| 888 | chunkwise_recurrents=[True, True, False, False], |
| 889 | layerscales=[False, False, False, False], layer_init_values=1e-6): |
| 890 | super().__init__() |
| 891 | |
| 892 | self.out_dim = out_dim |
| 893 | self.num_layers = len(depths) |
| 894 | self.embed_dim = embed_dims[0] |
| 895 | self.patch_norm = patch_norm |
| 896 | self.num_features = embed_dims[-1] |
| 897 | self.mlp_ratios = mlp_ratios |
| 898 | |
| 899 | # split image into non-overlapping patches |
| 900 | self.patch_embed = PatchEmbed(in_chans=in_chans, embed_dim=embed_dims[0], |
| 901 | norm_layer=norm_layer if self.patch_norm else None) |
| 902 | |
| 903 | # stochastic depth |
| 904 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule |
| 905 | |
| 906 | # build layers |
| 907 | self.layers = nn.ModuleList() |
| 908 | for i_layer in range(self.num_layers): |
| 909 | layer = BasicLayer( |
| 910 | embed_dim=embed_dims[i_layer], |
| 911 | out_dim=embed_dims[i_layer + 1] if (i_layer < self.num_layers - 1) else None, |
| 912 | depth=depths[i_layer], |
| 913 | num_heads=num_heads[i_layer], |
| 914 | init_value=init_values[i_layer], |
| 915 | heads_range=heads_ranges[i_layer], |
| 916 | mlp_ratio=mlp_ratios[i_layer], |
| 917 | split_size=split_sizes[i_layer], |
| 918 | sr_ratio=sr_ratios[i_layer], |
| 919 | # ffn_dim=int(mlp_ratios[i_layer] * embed_dims[i_layer]), |
| 920 | qkv_bias=True, |
| 921 | qk_scale=None, |
| 922 | drop_rate=0., |
| 923 | attn_drop=0.0, |
| 924 | drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], |
| 925 | # norm_layer=norm_layer, |
| 926 | chunkwise_recurrent=chunkwise_recurrents[i_layer], |
| 927 | downsample=PatchMerging if (i_layer in [0, 2]) else None, |
| 928 | use_checkpoint=use_checkpoints[i_layer], |
| 929 | mixer_type=mixer_types[i_layer], |
| 930 | layerscale=layerscales[i_layer], |
| 931 | layer_init_values=layer_init_values |
| 932 | ) |
| 933 | self.layers.append(layer) |
| 934 | |
| 935 | self.pooling = nn.AdaptiveAvgPool2d((embed_dims[self.num_layers - 1], 1)) |
| 936 | self.mlp_head = nn.Sequential( |
| 937 | nn.Linear(embed_dims[self.num_layers - 1], out_dim, bias=False), |
| 938 | nn.Hardswish(), |
| 939 | nn.Dropout(p=0.1) |
nothing calls this directly
no test coverage detected