(self, c_in=3, c_out=3, time_dim=256, num_classes=None, device="cuda")
| 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 / ( |
nothing calls this directly
no test coverage detected