| 79 | weight_init(self) |
| 80 | |
| 81 | class LayerNorm(nn.Module): |
| 82 | def __init__(self, dim, LayerNorm_type): |
| 83 | super(LayerNorm, self).__init__() |
| 84 | if LayerNorm_type =='BiasFree': |
| 85 | self.body = BiasFree_LayerNorm(dim) |
| 86 | else: |
| 87 | self.body = WithBias_LayerNorm(dim) |
| 88 | |
| 89 | def forward(self, x): |
| 90 | h, w = x.shape[-2:] |
| 91 | return to_4d(self.body(to_3d(x)), h, w) |
| 92 | |
| 93 | def initialize(self): |
| 94 | weight_init(self) |
| 95 | |
| 96 | class FeedForward(nn.Module): |
| 97 | def __init__(self, dim, ffn_expansion_factor, bias): |