(self)
| 19 | |
| 20 | class ImageMetricTracker(nn.Module): |
| 21 | def __init__(self): |
| 22 | super().__init__() |
| 23 | |
| 24 | self.ssim = SSIM(data_range=1.) |
| 25 | self.ssims = [] |
| 26 | |
| 27 | self.psnrs = [] |
| 28 | |
| 29 | self.fid = FrechetInceptionDistance( |
| 30 | feature=2048, |
| 31 | reset_real_features=True, |
| 32 | normalize=False, |
| 33 | sync_on_compute=True |
| 34 | ) |
| 35 | |
| 36 | def __call__(self, target, pred): |
| 37 | """ Assumes target and pred in [-1, 1] range """ |