| 237 | |
| 238 | |
| 239 | class Encoder(nn.Module): |
| 240 | def __init__(self, size, dim=512, dim_motion=20): |
| 241 | super(Encoder, self).__init__() |
| 242 | |
| 243 | # appearance netmork |
| 244 | self.net_app = EncoderApp(size, dim) |
| 245 | |
| 246 | # motion network |
| 247 | fc = [EqualLinear(dim, dim)] |
| 248 | for i in range(3): |
| 249 | fc.append(EqualLinear(dim, dim)) |
| 250 | |
| 251 | fc.append(EqualLinear(dim, dim_motion)) |
| 252 | self.fc = nn.Sequential(*fc) |
| 253 | |
| 254 | def enc_app(self, x): |
| 255 | |
| 256 | h_source = self.net_app(x) |
| 257 | |
| 258 | return h_source |
| 259 | |
| 260 | def enc_motion(self, x): |
| 261 | |
| 262 | h, _ = self.net_app(x) |
| 263 | h_motion = self.fc(h) |
| 264 | |
| 265 | return h_motion |
| 266 | |
| 267 | def forward(self, input_source, input_target, h_start=None): |
| 268 | |
| 269 | if input_target is not None: |
| 270 | |
| 271 | h_source, feats = self.net_app(input_source) |
| 272 | h_target, _ = self.net_app(input_target) |
| 273 | |
| 274 | h_motion_target = self.fc(h_target) |
| 275 | |
| 276 | if h_start is not None: |
| 277 | h_motion_source = self.fc(h_source) |
| 278 | h_motion = [h_motion_target, h_motion_source, h_start] |
| 279 | else: |
| 280 | h_motion = [h_motion_target] |
| 281 | |
| 282 | return h_source, h_motion, feats |
| 283 | else: |
| 284 | h_source, feats = self.net_app(input_source) |
| 285 | |
| 286 | return h_source, None, feats |