(self, in_channels, out_channels=None, dropout=0, up=False, num_groups=8, ks=3, input_norm=True, input_act=True)
| 107 | |
| 108 | class ResnetBlock(nn.Module): |
| 109 | def __init__(self, in_channels, out_channels=None, dropout=0, up=False, num_groups=8, ks=3, input_norm=True, input_act=True): |
| 110 | super().__init__() |
| 111 | self.in_channels = in_channels |
| 112 | out_channels = in_channels if out_channels is None else out_channels |
| 113 | self.out_channels = out_channels |
| 114 | self.up = up |
| 115 | |
| 116 | if input_norm and input_act: |
| 117 | self.in_layers = nn.Sequential( |
| 118 | nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True), |
| 119 | SiLU(), |
| 120 | nn.Conv2d(in_channels, out_channels, kernel_size=ks, stride=1, padding=(ks - 1)//2) |
| 121 | ) |
| 122 | elif not input_norm: |
| 123 | if input_act: |
| 124 | self.in_layers = nn.Sequential( |
| 125 | SiLU(), |
| 126 | nn.Conv2d(in_channels, out_channels, kernel_size=ks, stride=1, padding=(ks - 1)//2) |
| 127 | ) |
| 128 | else: |
| 129 | self.in_layers = nn.Sequential( |
| 130 | nn.Conv2d(in_channels, out_channels, kernel_size=ks, stride=1, padding=(ks - 1)//2) |
| 131 | ) |
| 132 | else: |
| 133 | raise NotImplementedError |
| 134 | |
| 135 | self.out_layers = nn.Sequential( |
| 136 | nn.GroupNorm(num_groups=num_groups, num_channels=out_channels, eps=1e-6, affine=True), |
| 137 | SiLU(), |
| 138 | nn.Dropout(p=dropout), |
| 139 | zero_module( |
| 140 | nn.Conv2d(out_channels, out_channels, kernel_size=ks, stride=1, padding=(ks - 1)//2) |
| 141 | ), |
| 142 | ) |
| 143 | |
| 144 | if self.in_channels != self.out_channels: |
| 145 | self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) |
| 146 | else: |
| 147 | self.shortcut = nn.Identity() |
| 148 | |
| 149 | def forward(self, x): |
| 150 | if self.up: |
nothing calls this directly
no test coverage detected