Spaces:
Runtime error
Runtime error
| # Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # | |
| # This work is made available under the Nvidia Source Code License-NC. | |
| # To view a copy of this license, check out LICENSE.md | |
| """ | |
| Modified from | |
| https://github.com/mseitzer/pytorch-fid | |
| Code adapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead | |
| of Tensorflow | |
| Copyright 2018 Institute of Bioinformatics, JKU Linz | |
| Licensed under the Apache License, Version 2.0 (the "License"); | |
| you may not use this file except in compliance with the License. | |
| You may obtain a copy of the License at | |
| http://www.apache.org/licenses/LICENSE-2.0 | |
| Unless required by applicable law or agreed to in writing, software | |
| distributed under the License is distributed on an "AS IS" BASIS, | |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| See the License for the specific language governing permissions and | |
| limitations under the License. | |
| """ | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| try: | |
| from torchvision.models.utils import load_state_dict_from_url | |
| except ImportError: | |
| from torch.utils.model_zoo import load_url as load_state_dict_from_url | |
| from torchvision.models import inception, inception_v3, vgg16 | |
| # Inception weights ported to Pytorch from | |
| # http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz | |
| FID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases' \ | |
| '/download/fid_weights/pt_inception-2015-12-05-6726825d.pth' | |
| class SwAV(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.model = torch.hub.load('facebookresearch/swav', 'resnet50', | |
| pretrained=True) | |
| self.model.fc = torch.nn.Sequential() | |
| def forward(self, x, align_corners=True): | |
| y = self.model(F.interpolate( | |
| x, size=(224, 224), mode='bicubic', align_corners=align_corners)) | |
| return y | |
| class Vgg16(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.model = vgg16(pretrained=True, init_weights=False) | |
| self.model.classifier = torch.nn.Sequential( | |
| *[self.model.classifier[i] for i in range(4)] | |
| ) | |
| def forward(self, x, align_corners=True): | |
| y = self.model(F.interpolate( | |
| x, size=(224, 224), mode='bicubic', align_corners=align_corners)) | |
| return y | |
| class InceptionV3(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.model = inception_v3(transform_input=False, | |
| pretrained=True, | |
| init_weights=False) | |
| self.model.fc = torch.nn.Sequential() | |
| def forward(self, x, align_corners=True): | |
| y = self.model(F.interpolate( | |
| x, size=(299, 299), mode='bicubic', align_corners=align_corners)) | |
| return y | |
| class TFInceptionV3(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.model = inception_v3(transform_input=False, | |
| num_classes=1008, | |
| aux_logits=False, | |
| pretrained=False, | |
| init_weights=False) | |
| self.model.Mixed_5b = FIDInceptionA(192, pool_features=32) | |
| self.model.Mixed_5c = FIDInceptionA(256, pool_features=64) | |
| self.model.Mixed_5d = FIDInceptionA(288, pool_features=64) | |
| self.model.Mixed_6b = FIDInceptionC(768, channels_7x7=128) | |
| self.model.Mixed_6c = FIDInceptionC(768, channels_7x7=160) | |
| self.model.Mixed_6d = FIDInceptionC(768, channels_7x7=160) | |
| self.model.Mixed_6e = FIDInceptionC(768, channels_7x7=192) | |
| self.model.Mixed_7b = FIDInceptionE_1(1280) | |
| self.model.Mixed_7c = FIDInceptionE_2(2048) | |
| state_dict = load_state_dict_from_url( | |
| FID_WEIGHTS_URL, progress=True, map_location='cpu' | |
| ) | |
| self.model.load_state_dict(state_dict) | |
| self.model.fc = torch.nn.Sequential() | |
| def forward(self, x, align_corners=True): | |
| y = self.model(F.interpolate( | |
| x, size=(299, 299), mode='bicubic', align_corners=align_corners)) | |
| return y | |
| class FIDInceptionA(inception.InceptionA): | |
| """InceptionA block patched for FID computation""" | |
| def __init__(self, in_channels, pool_features): | |
| super(FIDInceptionA, self).__init__(in_channels, pool_features) | |
| def forward(self, x): | |
| branch1x1 = self.branch1x1(x) | |
| branch5x5 = self.branch5x5_1(x) | |
| branch5x5 = self.branch5x5_2(branch5x5) | |
| branch3x3dbl = self.branch3x3dbl_1(x) | |
| branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) | |
| branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl) | |
| # Patch: Tensorflow's average pool does not use the padded zero's in | |
| # its average calculation | |
| branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, | |
| count_include_pad=False) | |
| branch_pool = self.branch_pool(branch_pool) | |
| outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool] | |
| return torch.cat(outputs, 1) | |
| class FIDInceptionC(inception.InceptionC): | |
| """InceptionC block patched for FID computation""" | |
| def __init__(self, in_channels, channels_7x7): | |
| super(FIDInceptionC, self).__init__(in_channels, channels_7x7) | |
| def forward(self, x): | |
| branch1x1 = self.branch1x1(x) | |
| branch7x7 = self.branch7x7_1(x) | |
| branch7x7 = self.branch7x7_2(branch7x7) | |
| branch7x7 = self.branch7x7_3(branch7x7) | |
| branch7x7dbl = self.branch7x7dbl_1(x) | |
| branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl) | |
| branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl) | |
| branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl) | |
| branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl) | |
| # Patch: Tensorflow's average pool does not use the padded zero's in | |
| # its average calculation | |
| branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, | |
| count_include_pad=False) | |
| branch_pool = self.branch_pool(branch_pool) | |
| outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool] | |
| return torch.cat(outputs, 1) | |
| class FIDInceptionE_1(inception.InceptionE): | |
| """First InceptionE block patched for FID computation""" | |
| def __init__(self, in_channels): | |
| super(FIDInceptionE_1, self).__init__(in_channels) | |
| def forward(self, x): | |
| branch1x1 = self.branch1x1(x) | |
| branch3x3 = self.branch3x3_1(x) | |
| branch3x3 = [ | |
| self.branch3x3_2a(branch3x3), | |
| self.branch3x3_2b(branch3x3), | |
| ] | |
| branch3x3 = torch.cat(branch3x3, 1) | |
| branch3x3dbl = self.branch3x3dbl_1(x) | |
| branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) | |
| branch3x3dbl = [ | |
| self.branch3x3dbl_3a(branch3x3dbl), | |
| self.branch3x3dbl_3b(branch3x3dbl), | |
| ] | |
| branch3x3dbl = torch.cat(branch3x3dbl, 1) | |
| # Patch: Tensorflow's average pool does not use the padded zero's in | |
| # its average calculation | |
| branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, | |
| count_include_pad=False) | |
| branch_pool = self.branch_pool(branch_pool) | |
| outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool] | |
| return torch.cat(outputs, 1) | |
| class FIDInceptionE_2(inception.InceptionE): | |
| """Second InceptionE block patched for FID computation""" | |
| def __init__(self, in_channels): | |
| super(FIDInceptionE_2, self).__init__(in_channels) | |
| def forward(self, x): | |
| branch1x1 = self.branch1x1(x) | |
| branch3x3 = self.branch3x3_1(x) | |
| branch3x3 = [ | |
| self.branch3x3_2a(branch3x3), | |
| self.branch3x3_2b(branch3x3), | |
| ] | |
| branch3x3 = torch.cat(branch3x3, 1) | |
| branch3x3dbl = self.branch3x3dbl_1(x) | |
| branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl) | |
| branch3x3dbl = [ | |
| self.branch3x3dbl_3a(branch3x3dbl), | |
| self.branch3x3dbl_3b(branch3x3dbl), | |
| ] | |
| branch3x3dbl = torch.cat(branch3x3dbl, 1) | |
| # Patch: The FID Inception model uses max pooling instead of average | |
| # pooling. This is likely an error in this specific Inception | |
| # implementation, as other Inception models use average pooling here | |
| # (which matches the description in the paper). | |
| branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1) | |
| branch_pool = self.branch_pool(branch_pool) | |
| outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool] | |
| return torch.cat(outputs, 1) | |