| 50 | |
| 51 | |
| 52 | class LPIPS(nn.Module): |
| 53 | def __init__(self): |
| 54 | super().__init__() |
| 55 | self.alexnet = AlexNet() |
| 56 | self.lpips_weights = nn.ModuleList() |
| 57 | for channels in self.alexnet.channels: |
| 58 | self.lpips_weights.append(Conv1x1(channels, 1)) |
| 59 | self._load_lpips_weights() |
| 60 | # imagenet normalization for range [-1, 1] |
| 61 | self.mu = torch.tensor([-0.03, -0.088, -0.188]).view(1, 3, 1, 1) # .cuda() |
| 62 | self.sigma = torch.tensor([0.458, 0.448, 0.450]).view(1, 3, 1, 1) # .cuda() |
| 63 | |
| 64 | def _load_lpips_weights(self): |
| 65 | own_state_dict = self.state_dict() |
| 66 | # if torch.cuda.is_available(): |
| 67 | # state_dict = torch.load('metrics/lpips_weights.ckpt') |
| 68 | # else: |
| 69 | state_dict = torch.load('lpips_weights.ckpt', |
| 70 | map_location=torch.device('cpu')) |
| 71 | for name, param in state_dict.items(): |
| 72 | if name in own_state_dict: |
| 73 | own_state_dict[name].copy_(param) |
| 74 | |
| 75 | def forward(self, x, y): |
| 76 | x = (x - self.mu.to(x.device)) / self.sigma.to(x.device) |
| 77 | y = (y - self.mu.to(x.device)) / self.sigma.to(x.device) |
| 78 | x_fmaps = self.alexnet(x) |
| 79 | y_fmaps = self.alexnet(y) |
| 80 | lpips_value = 0 |
| 81 | for x_fmap, y_fmap, conv1x1 in zip(x_fmaps, y_fmaps, self.lpips_weights): |
| 82 | x_fmap = normalize(x_fmap) |
| 83 | y_fmap = normalize(y_fmap) |
| 84 | lpips_value += torch.mean(conv1x1((x_fmap - y_fmap)**2)) |
| 85 | return lpips_value |
| 86 | |
| 87 | |
| 88 | @torch.no_grad() |
no outgoing calls
no test coverage detected