| 175 | |
| 176 | |
| 177 | class FaceEncoder(nn.Module): |
| 178 | def __init__(self, in_channel, out_dim): |
| 179 | super(FaceEncoder, self).__init__() |
| 180 | self.in_channel = in_channel |
| 181 | self.out_dim = out_dim |
| 182 | self.face_conv = nn.Sequential( |
| 183 | SameBlock2d(in_channel, 64, kernel_size=7, padding=3), |
| 184 | # # 64 → 32 |
| 185 | ResBlock2d(64, 64, kernel_size=3, padding=1), |
| 186 | DownBlock2d(64, 64, 3, 1), |
| 187 | SameBlock2d(64, 128), |
| 188 | # 32 → 16 |
| 189 | ResBlock2d(128, 128, kernel_size=3, padding=1), |
| 190 | DownBlock2d(128, 128, 3, 1), |
| 191 | SameBlock2d(128, 128), |
| 192 | # 16 → 8 |
| 193 | ResBlock2d(128, 128, kernel_size=3, padding=1), |
| 194 | DownBlock2d(128, 128, 3, 1), |
| 195 | SameBlock2d(128, 128), |
| 196 | # 8 → 4 |
| 197 | ResBlock2d(128, 128, kernel_size=3, padding=1), |
| 198 | DownBlock2d(128, 128, 3, 1), |
| 199 | SameBlock2d(128, 128), |
| 200 | # 4 → 2 |
| 201 | ResBlock2d(128, 128, kernel_size=3, padding=1), |
| 202 | DownBlock2d(128, 128, 3, 1), |
| 203 | SameBlock2d(128, out_dim, kernel_size=1, padding=0), |
| 204 | ) |
| 205 | |
| 206 | def forward(self, x): |
| 207 | ## b x c x h x w |
| 208 | out = self.face_conv(x) |
| 209 | return out |
| 210 | |
| 211 | |
| 212 | class AudioEncoder(nn.Module): |