Update the state with extracted features.
(self, videos: Tensor, real: bool)
| 370 | ) |
| 371 | |
| 372 | def update(self, videos: Tensor, real: bool) -> None: |
| 373 | """Update the state with extracted features.""" |
| 374 | videos = videos.permute(0, 4, 1, 2, 3) |
| 375 | |
| 376 | features = self.inception(videos).detach().cpu() |
| 377 | self.orig_dtype = features.dtype |
| 378 | features = features.double() |
| 379 | |
| 380 | if features.dim() == 1: |
| 381 | features = features.unsqueeze(0) |
| 382 | if real: |
| 383 | self.real_features_sum += features.sum(dim=0) |
| 384 | self.real_features_cov_sum += features.t().mm(features) |
| 385 | self.real_features_num_samples += videos.shape[0] |
| 386 | else: |
| 387 | self.fake_features_sum += features.sum(dim=0) |
| 388 | self.fake_features_cov_sum += features.t().mm(features) |
| 389 | self.fake_features_num_samples += videos.shape[0] |
| 390 | |
| 391 | def compute(self) -> Tensor: |
| 392 | """Calculate FID score based on accumulated extracted features from the two distributions.""" |