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