| 10 | |
| 11 | |
| 12 | class LPIPS(nn.Module): |
| 13 | # Learned perceptual metric |
| 14 | def __init__(self, use_dropout=True): |
| 15 | super().__init__() |
| 16 | self.scaling_layer = ScalingLayer() |
| 17 | self.chns = [64, 128, 256, 512, 512] # vg16 features |
| 18 | self.net = vgg16(pretrained=True, requires_grad=False) |
| 19 | self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout) |
| 20 | self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout) |
| 21 | self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout) |
| 22 | self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout) |
| 23 | self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout) |
| 24 | self.load_from_pretrained() |
| 25 | for param in self.parameters(): |
| 26 | param.requires_grad = False |
| 27 | |
| 28 | def load_from_pretrained(self, name="vgg_lpips"): |
| 29 | ckpt = get_ckpt_path(name, "sgm/modules/autoencoding/lpips/loss") |
| 30 | self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False) |
| 31 | print("loaded pretrained LPIPS loss from {}".format(ckpt)) |
| 32 | |
| 33 | @classmethod |
| 34 | def from_pretrained(cls, name="vgg_lpips"): |
| 35 | if name != "vgg_lpips": |
| 36 | raise NotImplementedError |
| 37 | model = cls() |
| 38 | ckpt = get_ckpt_path(name) |
| 39 | model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False) |
| 40 | return model |
| 41 | |
| 42 | def forward(self, input, target): |
| 43 | in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target)) |
| 44 | outs0, outs1 = self.net(in0_input), self.net(in1_input) |
| 45 | feats0, feats1, diffs = {}, {}, {} |
| 46 | lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4] |
| 47 | for kk in range(len(self.chns)): |
| 48 | feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk]) |
| 49 | diffs[kk] = (feats0[kk] - feats1[kk]) ** 2 |
| 50 | |
| 51 | res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))] |
| 52 | val = res[0] |
| 53 | for l in range(1, len(self.chns)): |
| 54 | val += res[l] |
| 55 | return val |
| 56 | |
| 57 | |
| 58 | class ScalingLayer(nn.Module): |