| 662 | |
| 663 | |
| 664 | class SimpleDecoder(nn.Module): |
| 665 | def __init__(self, in_channels, out_channels, *args, **kwargs): |
| 666 | super().__init__() |
| 667 | self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1), |
| 668 | ResnetBlock(in_channels=in_channels, |
| 669 | out_channels=2 * in_channels, |
| 670 | temb_channels=0, dropout=0.0), |
| 671 | ResnetBlock(in_channels=2 * in_channels, |
| 672 | out_channels=4 * in_channels, |
| 673 | temb_channels=0, dropout=0.0), |
| 674 | ResnetBlock(in_channels=4 * in_channels, |
| 675 | out_channels=2 * in_channels, |
| 676 | temb_channels=0, dropout=0.0), |
| 677 | nn.Conv2d(2*in_channels, in_channels, 1), |
| 678 | Upsample(in_channels, with_conv=True)]) |
| 679 | # end |
| 680 | self.norm_out = Normalize(in_channels) |
| 681 | self.conv_out = torch.nn.Conv2d(in_channels, |
| 682 | out_channels, |
| 683 | kernel_size=3, |
| 684 | stride=1, |
| 685 | padding=1) |
| 686 | |
| 687 | def forward(self, x): |
| 688 | for i, layer in enumerate(self.model): |
| 689 | if i in [1,2,3]: |
| 690 | x = layer(x, None) |
| 691 | else: |
| 692 | x = layer(x) |
| 693 | |
| 694 | h = self.norm_out(x) |
| 695 | h = nonlinearity(h) |
| 696 | x = self.conv_out(h) |
| 697 | return x |
| 698 | |
| 699 | |
| 700 | class UpsampleDecoder(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected