| 59 | class EdgeDecoderLayer(nn.Module): |
| 60 | |
| 61 | def __init__( |
| 62 | self, |
| 63 | dim, |
| 64 | num_heads, |
| 65 | mlp_ratio=4.0, |
| 66 | qkv_bias=False, |
| 67 | qk_scale=None, |
| 68 | drop=0.0, |
| 69 | attn_drop=0.0, |
| 70 | drop_path=[0.0, 0.0], |
| 71 | act_layer=nn.GELU, |
| 72 | norm_layer='nn.LayerNorm', |
| 73 | epsilon=1e-6, |
| 74 | ): |
| 75 | super().__init__() |
| 76 | |
| 77 | self.head_dim = dim // num_heads |
| 78 | self.scale = qk_scale or self.head_dim**-0.5 |
| 79 | |
| 80 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here |
| 81 | self.drop_path1 = DropPath( |
| 82 | drop_path[0]) if drop_path[0] > 0.0 else Identity() |
| 83 | self.norm1 = eval(norm_layer)(dim, epsilon=epsilon) |
| 84 | self.norm2 = eval(norm_layer)(dim, epsilon=epsilon) |
| 85 | |
| 86 | self.p = nn.Linear(dim, dim) |
| 87 | self.cv = nn.Linear(dim, dim) |
| 88 | self.pv = nn.Linear(dim, dim) |
| 89 | |
| 90 | self.dim = dim |
| 91 | self.num_heads = num_heads |
| 92 | self.p_proj = nn.Linear(dim, dim) |
| 93 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 94 | self.mlp_ratio = mlp_ratio |
| 95 | self.mlp = Mlp( |
| 96 | in_features=dim, |
| 97 | hidden_features=mlp_hidden_dim, |
| 98 | act_layer=act_layer, |
| 99 | drop=drop, |
| 100 | ) |
| 101 | |
| 102 | def forward(self, p, cv, pv): |
| 103 | pN = p.shape[1] |