(self, x: Tensor, pos=None,
num_patches=None, num_special=None, num_frames=None, enable_3d_rope=False)
| 75 | self.sample_drop_ratio = drop_path |
| 76 | |
| 77 | def forward(self, x: Tensor, pos=None, |
| 78 | num_patches=None, num_special=None, num_frames=None, enable_3d_rope=False) -> Tensor: |
| 79 | def attn_residual_func(x: Tensor, pos=None) -> Tensor: |
| 80 | return self.ls1(self.attn(self.norm1(x), pos=pos, |
| 81 | num_patches=num_patches, num_special=num_special, num_frames=num_frames, |
| 82 | enable_3d_rope=enable_3d_rope)) |
| 83 | |
| 84 | def ffn_residual_func(x: Tensor) -> Tensor: |
| 85 | return self.ls2(self.mlp(self.norm2(x))) |
| 86 | |
| 87 | if self.training and self.sample_drop_ratio > 0.1: |
| 88 | # the overhead is compensated only for a drop path rate larger than 0.1 |
| 89 | x = drop_add_residual_stochastic_depth( |
| 90 | x, pos=pos, residual_func=attn_residual_func, sample_drop_ratio=self.sample_drop_ratio |
| 91 | ) |
| 92 | x = drop_add_residual_stochastic_depth( |
| 93 | x, residual_func=ffn_residual_func, sample_drop_ratio=self.sample_drop_ratio |
| 94 | ) |
| 95 | elif self.training and self.sample_drop_ratio > 0.0: |
| 96 | x = x + self.drop_path1(attn_residual_func(x, pos=pos)) |
| 97 | x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 |
| 98 | else: |
| 99 | x = x + attn_residual_func(x, pos=pos) |
| 100 | x = x + ffn_residual_func(x) |
| 101 | return x |
| 102 | |
| 103 | |
| 104 | def drop_add_residual_stochastic_depth( |
nothing calls this directly
no test coverage detected