| 549 | init.kaiming_uniform_(m.weight) |
| 550 | |
| 551 | def forward_features(self, x): |
| 552 | x = self.patch_embed(x) |
| 553 | x = x + self.pos_embed |
| 554 | x = self.pos_drop(x) |
| 555 | for blk in self.blocks1: |
| 556 | x = blk(x) |
| 557 | if self.patch_merging is not None: |
| 558 | x = self.sub_sample1( |
| 559 | x.permute([0, 2, 1]).reshape( |
| 560 | [-1, self.embed_dim[0], self.HW[0], self.HW[1]])).contiguous() |
| 561 | for blk in self.blocks2: |
| 562 | x = blk(x) |
| 563 | if self.patch_merging is not None: |
| 564 | x = self.sub_sample2( |
| 565 | x.permute([0, 2, 1]).reshape( |
| 566 | [-1, self.embed_dim[1], self.HW[0] // 2, self.HW[1]])).contiguous() |
| 567 | for blk in self.blocks3: |
| 568 | x = blk(x) |
| 569 | if not self.prenorm: |
| 570 | x = self.norm(x) |
| 571 | return x |
| 572 | |
| 573 | def forward(self, x): |
| 574 | x = self.forward_features(x) |