| 160 | |
| 161 | |
| 162 | class DAT(nn.Module): |
| 163 | |
| 164 | def __init__(self, img_size=224, patch_size=4, num_classes=1000, expansion=4, |
| 165 | dim_stem=96, dims=[96, 192, 384, 768], depths=[2, 2, 6, 2], |
| 166 | heads=[3, 6, 12, 24], heads_q=[6, 12, 24, 48], |
| 167 | window_sizes=[7, 7, 7, 7], |
| 168 | drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0, |
| 169 | strides=[-1,-1,-1,-1], |
| 170 | offset_range_factor=[1, 2, 3, 4], |
| 171 | stage_spec=[['L', 'D'], ['L', 'D'], ['L', 'D', 'L', 'D', 'L', 'D'], ['L', 'D']], |
| 172 | groups=[-1, -1, 3, 6], |
| 173 | use_pes=[False, False, False, False], |
| 174 | dwc_pes=[False, False, False, False], |
| 175 | sr_ratios=[8, 4, 2, 1], |
| 176 | lower_lr_kvs={}, |
| 177 | fixed_pes=[False, False, False, False], |
| 178 | no_offs=[False, False, False, False], |
| 179 | ns_per_pts=[4, 4, 4, 4], |
| 180 | use_dwc_mlps=[False, False, False, False], |
| 181 | use_conv_patches=False, |
| 182 | ksizes=[9, 7, 5, 3], |
| 183 | ksize_qnas=[3, 3, 3, 3], |
| 184 | nqs=[2, 2, 2, 2], |
| 185 | qna_activation='exp', |
| 186 | nat_ksizes=[3,3,3,3], |
| 187 | layer_scale_values=[-1,-1,-1,-1], |
| 188 | use_lpus=[False, False, False, False], |
| 189 | log_cpb=[False, False, False, False], |
| 190 | **kwargs): |
| 191 | super().__init__() |
| 192 | |
| 193 | self.patch_proj = nn.Sequential( |
| 194 | nn.Conv2d(3, dim_stem // 2, 3, patch_size // 2, 1), |
| 195 | LayerNormProxy(dim_stem // 2), |
| 196 | nn.GELU(), |
| 197 | nn.Conv2d(dim_stem // 2, dim_stem, 3, patch_size // 2, 1), |
| 198 | LayerNormProxy(dim_stem) |
| 199 | ) if use_conv_patches else nn.Sequential( |
| 200 | nn.Conv2d(3, dim_stem, patch_size, patch_size, 0), |
| 201 | LayerNormProxy(dim_stem) |
| 202 | ) |
| 203 | |
| 204 | img_size = img_size // patch_size |
| 205 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] |
| 206 | |
| 207 | self.stages = nn.ModuleList() |
| 208 | for i in range(4): |
| 209 | dim1 = dim_stem if i == 0 else dims[i - 1] * 2 |
| 210 | dim2 = dims[i] |
| 211 | self.stages.append( |
| 212 | TransformerStage( |
| 213 | img_size, window_sizes[i], ns_per_pts[i], |
| 214 | dim1, dim2, depths[i], |
| 215 | stage_spec[i], groups[i], use_pes[i], |
| 216 | sr_ratios[i], heads[i], heads_q[i], strides[i], |
| 217 | offset_range_factor[i], |
| 218 | dwc_pes[i], no_offs[i], fixed_pes[i], |
| 219 | attn_drop_rate, drop_rate, expansion, drop_rate, |