Learned perceptual metric.
| 59 | |
| 60 | |
| 61 | class LPIPS(nn.Module): |
| 62 | """Learned perceptual metric.""" |
| 63 | |
| 64 | def __init__(self, use_dropout=True): |
| 65 | super().__init__() |
| 66 | self.scaling_layer = ScalingLayer() |
| 67 | self.chns = [64, 128, 256, 512, 512] # vgg16 features |
| 68 | self.net = vgg16(pretrained=True, requires_grad=False) |
| 69 | self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout) |
| 70 | self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout) |
| 71 | self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout) |
| 72 | self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout) |
| 73 | self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout) |
| 74 | self.load_from_pretrained() |
| 75 | for param in self.parameters(): |
| 76 | param.requires_grad = False |
| 77 | |
| 78 | def load_from_pretrained(self, name="vgg_lpips"): |
| 79 | ckpt = get_ckpt_path(name) |
| 80 | self.load_state_dict( |
| 81 | torch.load(ckpt, map_location=torch.device("cpu")), strict=False |
| 82 | ) |
| 83 | |
| 84 | def forward(self, input, target): |
| 85 | in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target)) |
| 86 | outs0, outs1 = self.net(in0_input), self.net(in1_input) |
| 87 | feats0, feats1, diffs = {}, {}, {} |
| 88 | lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4] |
| 89 | for kk in range(len(self.chns)): |
| 90 | feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk]) |
| 91 | diffs[kk] = (feats0[kk] - feats1[kk]) ** 2 |
| 92 | |
| 93 | res = [ |
| 94 | spatial_average(lins[kk].model(diffs[kk]), keepdim=True) |
| 95 | for kk in range(len(self.chns)) |
| 96 | ] |
| 97 | val = res[0] |
| 98 | for l in range(1, len(self.chns)): |
| 99 | val += res[l] |
| 100 | return val |
| 101 | |
| 102 | |
| 103 | class ScalingLayer(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected