| 187 | |
| 188 | |
| 189 | class UNet_conditional(nn.Module): |
| 190 | def __init__(self, c_in=3, c_out=3, time_dim=256, num_classes=None, device="cuda"): |
| 191 | super().__init__() |
| 192 | self.device = device |
| 193 | self.time_dim = time_dim |
| 194 | self.inc = DoubleConv(c_in, 64) |
| 195 | self.down1 = Down(64, 128) |
| 196 | self.sa1 = SelfAttention(128, 32) |
| 197 | self.down2 = Down(128, 256) |
| 198 | self.sa2 = SelfAttention(256, 16) |
| 199 | self.down3 = Down(256, 256) |
| 200 | self.sa3 = SelfAttention(256, 8) |
| 201 | |
| 202 | self.bot1 = DoubleConv(256, 512) |
| 203 | self.bot2 = DoubleConv(512, 512) |
| 204 | self.bot3 = DoubleConv(512, 256) |
| 205 | |
| 206 | self.up1 = Up(512, 128) |
| 207 | self.sa4 = SelfAttention(128, 16) |
| 208 | self.up2 = Up(256, 64) |
| 209 | self.sa5 = SelfAttention(64, 32) |
| 210 | self.up3 = Up(128, 64) |
| 211 | self.sa6 = SelfAttention(64, 64) |
| 212 | self.outc = nn.Conv2d(64, c_out, kernel_size=1) |
| 213 | |
| 214 | if num_classes is not None: |
| 215 | self.label_emb = nn.Embedding(num_classes, time_dim) |
| 216 | |
| 217 | def pos_encoding(self, t, channels): |
| 218 | inv_freq = 1.0 / ( |
| 219 | 10000 |
| 220 | ** (torch.arange(0, channels, 2, device=self.device).float() / channels) |
| 221 | ) |
| 222 | pos_enc_a = torch.sin(t.repeat(1, channels // 2) * inv_freq) |
| 223 | pos_enc_b = torch.cos(t.repeat(1, channels // 2) * inv_freq) |
| 224 | pos_enc = torch.cat([pos_enc_a, pos_enc_b], dim=-1) |
| 225 | return pos_enc |
| 226 | |
| 227 | def forward(self, x, t, y): |
| 228 | t = t.unsqueeze(-1).type(torch.float) |
| 229 | t = self.pos_encoding(t, self.time_dim) |
| 230 | |
| 231 | if y is not None: |
| 232 | t += self.label_emb(y) |
| 233 | |
| 234 | x1 = self.inc(x) |
| 235 | x2 = self.down1(x1, t) |
| 236 | x2 = self.sa1(x2) |
| 237 | x3 = self.down2(x2, t) |
| 238 | x3 = self.sa2(x3) |
| 239 | x4 = self.down3(x3, t) |
| 240 | x4 = self.sa3(x4) |
| 241 | |
| 242 | x4 = self.bot1(x4) |
| 243 | x4 = self.bot2(x4) |
| 244 | x4 = self.bot3(x4) |
| 245 | |
| 246 | x = self.up1(x4, x3, t) |
no outgoing calls
no test coverage detected