(self, c_in, n_head, norm, dropout_rate=0.)
| 317 | |
| 318 | class SelfAttention2d(ConditionedModule): |
| 319 | def __init__(self, c_in, n_head, norm, dropout_rate=0.): |
| 320 | super().__init__() |
| 321 | assert c_in % n_head == 0 |
| 322 | self.norm_in = norm(c_in) |
| 323 | self.n_head = n_head |
| 324 | self.qkv_proj = nn.Conv2d(c_in, c_in * 3, 1) |
| 325 | self.out_proj = nn.Conv2d(c_in, c_in, 1) |
| 326 | self.dropout = nn.Dropout(dropout_rate) |
| 327 | nn.init.zeros_(self.out_proj.weight) |
| 328 | nn.init.zeros_(self.out_proj.bias) |
| 329 | |
| 330 | def forward(self, input, cond): |
| 331 | n, c, h, w = input.shape |