commit

LICENSE

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2019 Kim Seonghyeon
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

LICENSE-FID

@@ -1,201 +0,0 @@
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright [yyyy] [name of copyright owner]
-
-   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.

LICENSE-LPIPS

@@ -1,24 +0,0 @@
-Copyright (c) 2018, Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, Oliver Wang
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-* Redistributions of source code must retain the above copyright notice, this
-  list of conditions and the following disclaimer.
-
-* Redistributions in binary form must reproduce the above copyright notice,
-  this list of conditions and the following disclaimer in the documentation
-  and/or other materials provided with the distribution.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-

LICENSE-NVIDIA

@@ -1,101 +0,0 @@
-Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
-
-
-Nvidia Source Code License-NC
-
-=======================================================================
-
-1. Definitions
-
-"Licensor" means any person or entity that distributes its Work.
-
-"Software" means the original work of authorship made available under
-this License.
-
-"Work" means the Software and any additions to or derivative works of
-the Software that are made available under this License.
-
-"Nvidia Processors" means any central processing unit (CPU), graphics
-processing unit (GPU), field-programmable gate array (FPGA),
-application-specific integrated circuit (ASIC) or any combination
-thereof designed, made, sold, or provided by Nvidia or its affiliates.
-
-The terms "reproduce," "reproduction," "derivative works," and
-"distribution" have the meaning as provided under U.S. copyright law;
-provided, however, that for the purposes of this License, derivative
-works shall not include works that remain separable from, or merely
-link (or bind by name) to the interfaces of, the Work.
-
-Works, including the Software, are "made available" under this License
-by including in or with the Work either (a) a copyright notice
-referencing the applicability of this License to the Work, or (b) a
-copy of this License.
-
-2. License Grants
-
-    2.1 Copyright Grant. Subject to the terms and conditions of this
-    License, each Licensor grants to you a perpetual, worldwide,
-    non-exclusive, royalty-free, copyright license to reproduce,
-    prepare derivative works of, publicly display, publicly perform,
-    sublicense and distribute its Work and any resulting derivative
-    works in any form.
-
-3. Limitations
-
-    3.1 Redistribution. You may reproduce or distribute the Work only
-    if (a) you do so under this License, (b) you include a complete
-    copy of this License with your distribution, and (c) you retain
-    without modification any copyright, patent, trademark, or
-    attribution notices that are present in the Work.
-
-    3.2 Derivative Works. You may specify that additional or different
-    terms apply to the use, reproduction, and distribution of your
-    derivative works of the Work ("Your Terms") only if (a) Your Terms
-    provide that the use limitation in Section 3.3 applies to your
-    derivative works, and (b) you identify the specific derivative
-    works that are subject to Your Terms. Notwithstanding Your Terms,
-    this License (including the redistribution requirements in Section
-    3.1) will continue to apply to the Work itself.
-
-    3.3 Use Limitation. The Work and any derivative works thereof only
-    may be used or intended for use non-commercially. The Work or
-    derivative works thereof may be used or intended for use by Nvidia
-    or its affiliates commercially or non-commercially. As used herein,
-    "non-commercially" means for research or evaluation purposes only.
-
-    3.4 Patent Claims. If you bring or threaten to bring a patent claim
-    against any Licensor (including any claim, cross-claim or
-    counterclaim in a lawsuit) to enforce any patents that you allege
-    are infringed by any Work, then your rights under this License from
-    such Licensor (including the grants in Sections 2.1 and 2.2) will
-    terminate immediately.
-
-    3.5 Trademarks. This License does not grant any rights to use any
-    Licensor's or its affiliates' names, logos, or trademarks, except
-    as necessary to reproduce the notices described in this License.
-
-    3.6 Termination. If you violate any term of this License, then your
-    rights under this License (including the grants in Sections 2.1 and
-    2.2) will terminate immediately.
-
-4. Disclaimer of Warranty.
-
-THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
-KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
-MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
-NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
-THIS LICENSE. 
-
-5. Limitation of Liability.
-
-EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
-THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
-SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
-INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
-OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
-(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
-LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
-COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
-THE POSSIBILITY OF SUCH DAMAGES.
-
-=======================================================================

apply_factor.py

@@ -1,94 +0,0 @@
-import argparse
-
-import torch
-from torchvision import utils
-
-from model import Generator
-
-
-if __name__ == "__main__":
-    torch.set_grad_enabled(False)
-
-    parser = argparse.ArgumentParser(description="Apply closed form factorization")
-
-    parser.add_argument(
-        "-i", "--index", type=int, default=0, help="index of eigenvector"
-    )
-    parser.add_argument(
-        "-d",
-        "--degree",
-        type=float,
-        default=5,
-        help="scalar factors for moving latent vectors along eigenvector",
-    )
-    parser.add_argument(
-        "--channel_multiplier",
-        type=int,
-        default=2,
-        help='channel multiplier factor. config-f = 2, else = 1',
-    )
-    parser.add_argument("--ckpt", type=str, required=True, help="stylegan2 checkpoints")
-    parser.add_argument(
-        "--size", type=int, default=256, help="output image size of the generator"
-    )
-    parser.add_argument(
-        "-n", "--n_sample", type=int, default=7, help="number of samples created"
-    )
-    parser.add_argument(
-        "--truncation", type=float, default=0.7, help="truncation factor"
-    )
-    parser.add_argument(
-        "--device", type=str, default="cuda", help="device to run the model"
-    )
-    parser.add_argument(
-        "--out_prefix",
-        type=str,
-        default="factor",
-        help="filename prefix to result samples",
-    )
-    parser.add_argument(
-        "factor",
-        type=str,
-        help="name of the closed form factorization result factor file",
-    )
-
-    args = parser.parse_args()
-
-    eigvec = torch.load(args.factor)["eigvec"].to(args.device)
-    ckpt = torch.load(args.ckpt)
-    g = Generator(args.size, 512, 8, channel_multiplier=args.channel_multiplier).to(args.device)
-    g.load_state_dict(ckpt["g_ema"], strict=False)
-
-    trunc = g.mean_latent(4096)
-
-    latent = torch.randn(args.n_sample, 512, device=args.device)
-    latent = g.get_latent(latent)
-
-    direction = args.degree * eigvec[:, args.index].unsqueeze(0)
-
-    img, _ = g(
-        [latent],
-        truncation=args.truncation,
-        truncation_latent=trunc,
-        input_is_latent=True,
-    )
-    img1, _ = g(
-        [latent + direction],
-        truncation=args.truncation,
-        truncation_latent=trunc,
-        input_is_latent=True,
-    )
-    img2, _ = g(
-        [latent - direction],
-        truncation=args.truncation,
-        truncation_latent=trunc,
-        input_is_latent=True,
-    )
-
-    grid = utils.save_image(
-        torch.cat([img1, img, img2], 0),
-        f"{args.out_prefix}_index-{args.index}_degree-{args.degree}.png",
-        normalize=True,
-        range=(-1, 1),
-        nrow=args.n_sample,
-    )

calc_inception.py

@@ -1,130 +0,0 @@
-import argparse
-import pickle
-import os
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.utils.data import DataLoader
-from torchvision import transforms
-from torchvision.models import inception_v3, Inception3
-import numpy as np
-from tqdm import tqdm
-
-from inception import InceptionV3
-from dataset import MultiResolutionDataset
-
-
-class Inception3Feature(Inception3):
-    def forward(self, x):
-        if x.shape[2] != 299 or x.shape[3] != 299:
-            x = F.interpolate(x, size=(299, 299), mode="bilinear", align_corners=True)
-
-        x = self.Conv2d_1a_3x3(x)  # 299 x 299 x 3
-        x = self.Conv2d_2a_3x3(x)  # 149 x 149 x 32
-        x = self.Conv2d_2b_3x3(x)  # 147 x 147 x 32
-        x = F.max_pool2d(x, kernel_size=3, stride=2)  # 147 x 147 x 64
-
-        x = self.Conv2d_3b_1x1(x)  # 73 x 73 x 64
-        x = self.Conv2d_4a_3x3(x)  # 73 x 73 x 80
-        x = F.max_pool2d(x, kernel_size=3, stride=2)  # 71 x 71 x 192
-
-        x = self.Mixed_5b(x)  # 35 x 35 x 192
-        x = self.Mixed_5c(x)  # 35 x 35 x 256
-        x = self.Mixed_5d(x)  # 35 x 35 x 288
-
-        x = self.Mixed_6a(x)  # 35 x 35 x 288
-        x = self.Mixed_6b(x)  # 17 x 17 x 768
-        x = self.Mixed_6c(x)  # 17 x 17 x 768
-        x = self.Mixed_6d(x)  # 17 x 17 x 768
-        x = self.Mixed_6e(x)  # 17 x 17 x 768
-
-        x = self.Mixed_7a(x)  # 17 x 17 x 768
-        x = self.Mixed_7b(x)  # 8 x 8 x 1280
-        x = self.Mixed_7c(x)  # 8 x 8 x 2048
-
-        x = F.avg_pool2d(x, kernel_size=8)  # 8 x 8 x 2048
-
-        return x.view(x.shape[0], x.shape[1])  # 1 x 1 x 2048
-
-
-def load_patched_inception_v3():
-    # inception = inception_v3(pretrained=True)
-    # inception_feat = Inception3Feature()
-    # inception_feat.load_state_dict(inception.state_dict())
-    inception_feat = InceptionV3([3], normalize_input=False)
-
-    return inception_feat
-
-
-@torch.no_grad()
-def extract_features(loader, inception, device):
-    pbar = tqdm(loader)
-
-    feature_list = []
-
-    for img in pbar:
-        img = img.to(device)
-        feature = inception(img)[0].view(img.shape[0], -1)
-        feature_list.append(feature.to("cpu"))
-
-    features = torch.cat(feature_list, 0)
-
-    return features
-
-
-if __name__ == "__main__":
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    parser = argparse.ArgumentParser(
-        description="Calculate Inception v3 features for datasets"
-    )
-    parser.add_argument(
-        "--size",
-        type=int,
-        default=256,
-        help="image sizes used for embedding calculation",
-    )
-    parser.add_argument(
-        "--batch", default=64, type=int, help="batch size for inception networks"
-    )
-    parser.add_argument(
-        "--n_sample",
-        type=int,
-        default=50000,
-        help="number of samples used for embedding calculation",
-    )
-    parser.add_argument(
-        "--flip", action="store_true", help="apply random flipping to real images"
-    )
-    parser.add_argument("path", metavar="PATH", help="path to datset lmdb file")
-
-    args = parser.parse_args()
-
-    inception = load_patched_inception_v3()
-    inception = nn.DataParallel(inception).eval().to(device)
-
-    transform = transforms.Compose(
-        [
-            transforms.RandomHorizontalFlip(p=0.5 if args.flip else 0),
-            transforms.ToTensor(),
-            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
-        ]
-    )
-
-    dset = MultiResolutionDataset(args.path, transform=transform, resolution=args.size)
-    loader = DataLoader(dset, batch_size=args.batch, num_workers=4)
-
-    features = extract_features(loader, inception, device).numpy()
-
-    features = features[: args.n_sample]
-
-    print(f"extracted {features.shape[0]} features")
-
-    mean = np.mean(features, 0)
-    cov = np.cov(features, rowvar=False)
-
-    name = os.path.splitext(os.path.basename(args.path))[0]
-
-    with open(f"inception_{name}.pkl", "wb") as f:
-        pickle.dump({"mean": mean, "cov": cov, "size": args.size, "path": args.path}, f)

checkpoint/.gitignore

@@ -1 +0,0 @@
-*.pt

closed_form_factorization.py

@@ -1,33 +0,0 @@
-import argparse
-
-import torch
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Extract factor/eigenvectors of latent spaces using closed form factorization"
-    )
-
-    parser.add_argument(
-        "--out", type=str, default="factor.pt", help="name of the result factor file"
-    )
-    parser.add_argument("ckpt", type=str, help="name of the model checkpoint")
-
-    args = parser.parse_args()
-
-    ckpt = torch.load(args.ckpt)
-    modulate = {
-        k: v
-        for k, v in ckpt["g_ema"].items()
-        if "modulation" in k and "to_rgbs" not in k and "weight" in k
-    }
-
-    weight_mat = []
-    for k, v in modulate.items():
-        weight_mat.append(v)
-
-    W = torch.cat(weight_mat, 0)
-    eigvec = torch.svd(W).V.to("cpu")
-
-    torch.save({"ckpt": args.ckpt, "eigvec": eigvec}, args.out)
-

convert_weight.py

@@ -1,301 +0,0 @@
-import argparse
-import os
-import sys
-import pickle
-import math
-
-import torch
-import numpy as np
-from torchvision import utils
-
-from model import Generator, Discriminator
-
-
-def convert_modconv(vars, source_name, target_name, flip=False):
-    weight = vars[source_name + "/weight"].value().eval()
-    mod_weight = vars[source_name + "/mod_weight"].value().eval()
-    mod_bias = vars[source_name + "/mod_bias"].value().eval()
-    noise = vars[source_name + "/noise_strength"].value().eval()
-    bias = vars[source_name + "/bias"].value().eval()
-
-    dic = {
-        "conv.weight": np.expand_dims(weight.transpose((3, 2, 0, 1)), 0),
-        "conv.modulation.weight": mod_weight.transpose((1, 0)),
-        "conv.modulation.bias": mod_bias + 1,
-        "noise.weight": np.array([noise]),
-        "activate.bias": bias,
-    }
-
-    dic_torch = {}
-
-    for k, v in dic.items():
-        dic_torch[target_name + "." + k] = torch.from_numpy(v)
-
-    if flip:
-        dic_torch[target_name + ".conv.weight"] = torch.flip(
-            dic_torch[target_name + ".conv.weight"], [3, 4]
-        )
-
-    return dic_torch
-
-
-def convert_conv(vars, source_name, target_name, bias=True, start=0):
-    weight = vars[source_name + "/weight"].value().eval()
-
-    dic = {"weight": weight.transpose((3, 2, 0, 1))}
-
-    if bias:
-        dic["bias"] = vars[source_name + "/bias"].value().eval()
-
-    dic_torch = {}
-
-    dic_torch[target_name + f".{start}.weight"] = torch.from_numpy(dic["weight"])
-
-    if bias:
-        dic_torch[target_name + f".{start + 1}.bias"] = torch.from_numpy(dic["bias"])
-
-    return dic_torch
-
-
-def convert_torgb(vars, source_name, target_name):
-    weight = vars[source_name + "/weight"].value().eval()
-    mod_weight = vars[source_name + "/mod_weight"].value().eval()
-    mod_bias = vars[source_name + "/mod_bias"].value().eval()
-    bias = vars[source_name + "/bias"].value().eval()
-
-    dic = {
-        "conv.weight": np.expand_dims(weight.transpose((3, 2, 0, 1)), 0),
-        "conv.modulation.weight": mod_weight.transpose((1, 0)),
-        "conv.modulation.bias": mod_bias + 1,
-        "bias": bias.reshape((1, 3, 1, 1)),
-    }
-
-    dic_torch = {}
-
-    for k, v in dic.items():
-        dic_torch[target_name + "." + k] = torch.from_numpy(v)
-
-    return dic_torch
-
-
-def convert_dense(vars, source_name, target_name):
-    weight = vars[source_name + "/weight"].value().eval()
-    bias = vars[source_name + "/bias"].value().eval()
-
-    dic = {"weight": weight.transpose((1, 0)), "bias": bias}
-
-    dic_torch = {}
-
-    for k, v in dic.items():
-        dic_torch[target_name + "." + k] = torch.from_numpy(v)
-
-    return dic_torch
-
-
-def update(state_dict, new):
-    for k, v in new.items():
-        if k not in state_dict:
-            raise KeyError(k + " is not found")
-
-        if v.shape != state_dict[k].shape:
-            raise ValueError(f"Shape mismatch: {v.shape} vs {state_dict[k].shape}")
-
-        state_dict[k] = v
-
-
-def discriminator_fill_statedict(statedict, vars, size):
-    log_size = int(math.log(size, 2))
-
-    update(statedict, convert_conv(vars, f"{size}x{size}/FromRGB", "convs.0"))
-
-    conv_i = 1
-
-    for i in range(log_size - 2, 0, -1):
-        reso = 4 * 2 ** i
-        update(
-            statedict,
-            convert_conv(vars, f"{reso}x{reso}/Conv0", f"convs.{conv_i}.conv1"),
-        )
-        update(
-            statedict,
-            convert_conv(
-                vars, f"{reso}x{reso}/Conv1_down", f"convs.{conv_i}.conv2", start=1
-            ),
-        )
-        update(
-            statedict,
-            convert_conv(
-                vars, f"{reso}x{reso}/Skip", f"convs.{conv_i}.skip", start=1, bias=False
-            ),
-        )
-        conv_i += 1
-
-    update(statedict, convert_conv(vars, f"4x4/Conv", "final_conv"))
-    update(statedict, convert_dense(vars, f"4x4/Dense0", "final_linear.0"))
-    update(statedict, convert_dense(vars, f"Output", "final_linear.1"))
-
-    return statedict
-
-
-def fill_statedict(state_dict, vars, size, n_mlp):
-    log_size = int(math.log(size, 2))
-
-    for i in range(n_mlp):
-        update(state_dict, convert_dense(vars, f"G_mapping/Dense{i}", f"style.{i + 1}"))
-
-    update(
-        state_dict,
-        {
-            "input.input": torch.from_numpy(
-                vars["G_synthesis/4x4/Const/const"].value().eval()
-            )
-        },
-    )
-
-    update(state_dict, convert_torgb(vars, "G_synthesis/4x4/ToRGB", "to_rgb1"))
-
-    for i in range(log_size - 2):
-        reso = 4 * 2 ** (i + 1)
-        update(
-            state_dict,
-            convert_torgb(vars, f"G_synthesis/{reso}x{reso}/ToRGB", f"to_rgbs.{i}"),
-        )
-
-    update(state_dict, convert_modconv(vars, "G_synthesis/4x4/Conv", "conv1"))
-
-    conv_i = 0
-
-    for i in range(log_size - 2):
-        reso = 4 * 2 ** (i + 1)
-        update(
-            state_dict,
-            convert_modconv(
-                vars,
-                f"G_synthesis/{reso}x{reso}/Conv0_up",
-                f"convs.{conv_i}",
-                flip=True,
-            ),
-        )
-        update(
-            state_dict,
-            convert_modconv(
-                vars, f"G_synthesis/{reso}x{reso}/Conv1", f"convs.{conv_i + 1}"
-            ),
-        )
-        conv_i += 2
-
-    for i in range(0, (log_size - 2) * 2 + 1):
-        update(
-            state_dict,
-            {
-                f"noises.noise_{i}": torch.from_numpy(
-                    vars[f"G_synthesis/noise{i}"].value().eval()
-                )
-            },
-        )
-
-    return state_dict
-
-
-if __name__ == "__main__":
-    device = "cuda"
-
-    parser = argparse.ArgumentParser(
-        description="Tensorflow to pytorch model checkpoint converter"
-    )
-    parser.add_argument(
-        "--repo",
-        type=str,
-        required=True,
-        help="path to the offical StyleGAN2 repository with dnnlib/ folder",
-    )
-    parser.add_argument(
-        "--gen", action="store_true", help="convert the generator weights"
-    )
-    parser.add_argument(
-        "--disc", action="store_true", help="convert the discriminator weights"
-    )
-    parser.add_argument(
-        "--channel_multiplier",
-        type=int,
-        default=2,
-        help="channel multiplier factor. config-f = 2, else = 1",
-    )
-    parser.add_argument("path", metavar="PATH", help="path to the tensorflow weights")
-
-    args = parser.parse_args()
-
-    sys.path.append(args.repo)
-
-    import dnnlib
-    from dnnlib import tflib
-
-    tflib.init_tf()
-
-    with open(args.path, "rb") as f:
-        generator, discriminator, g_ema = pickle.load(f)
-
-    size = g_ema.output_shape[2]
-
-    n_mlp = 0
-    mapping_layers_names = g_ema.__getstate__()['components']['mapping'].list_layers()
-    for layer in mapping_layers_names:
-        if layer[0].startswith('Dense'):
-            n_mlp += 1
-
-    g = Generator(size, 512, n_mlp, channel_multiplier=args.channel_multiplier)
-    state_dict = g.state_dict()
-    state_dict = fill_statedict(state_dict, g_ema.vars, size, n_mlp)
-
-    g.load_state_dict(state_dict)
-
-    latent_avg = torch.from_numpy(g_ema.vars["dlatent_avg"].value().eval())
-
-    ckpt = {"g_ema": state_dict, "latent_avg": latent_avg}
-
-    if args.gen:
-        g_train = Generator(size, 512, n_mlp, channel_multiplier=args.channel_multiplier)
-        g_train_state = g_train.state_dict()
-        g_train_state = fill_statedict(g_train_state, generator.vars, size, n_mlp)
-        ckpt["g"] = g_train_state
-
-    if args.disc:
-        disc = Discriminator(size, channel_multiplier=args.channel_multiplier)
-        d_state = disc.state_dict()
-        d_state = discriminator_fill_statedict(d_state, discriminator.vars, size)
-        ckpt["d"] = d_state
-
-    name = os.path.splitext(os.path.basename(args.path))[0]
-    torch.save(ckpt, name + ".pt")
-
-    batch_size = {256: 16, 512: 9, 1024: 4}
-    n_sample = batch_size.get(size, 25)
-
-    g = g.to(device)
-
-    z = np.random.RandomState(0).randn(n_sample, 512).astype("float32")
-
-    with torch.no_grad():
-        img_pt, _ = g(
-            [torch.from_numpy(z).to(device)],
-            truncation=0.5,
-            truncation_latent=latent_avg.to(device),
-            randomize_noise=False,
-        )
-
-    Gs_kwargs = dnnlib.EasyDict()
-    Gs_kwargs.randomize_noise = False
-    img_tf = g_ema.run(z, None, **Gs_kwargs)
-    img_tf = torch.from_numpy(img_tf).to(device)
-
-    img_diff = ((img_pt + 1) / 2).clamp(0.0, 1.0) - ((img_tf.to(device) + 1) / 2).clamp(
-        0.0, 1.0
-    )
-
-    img_concat = torch.cat((img_tf, img_pt, img_diff), dim=0)
-
-    print(img_diff.abs().max())
-
-    utils.save_image(
-        img_concat, name + ".png", nrow=n_sample, normalize=True, range=(-1, 1)
-    )

dataset.py

@@ -1,40 +0,0 @@
-from io import BytesIO
-
-import lmdb
-from PIL import Image
-from torch.utils.data import Dataset
-
-
-class MultiResolutionDataset(Dataset):
-    def __init__(self, path, transform, resolution=256):
-        self.env = lmdb.open(
-            path,
-            max_readers=32,
-            readonly=True,
-            lock=False,
-            readahead=False,
-            meminit=False,
-        )
-
-        if not self.env:
-            raise IOError('Cannot open lmdb dataset', path)
-
-        with self.env.begin(write=False) as txn:
-            self.length = int(txn.get('length'.encode('utf-8')).decode('utf-8'))
-
-        self.resolution = resolution
-        self.transform = transform
-
-    def __len__(self):
-        return self.length
-
-    def __getitem__(self, index):
-        with self.env.begin(write=False) as txn:
-            key = f'{self.resolution}-{str(index).zfill(5)}'.encode('utf-8')
-            img_bytes = txn.get(key)
-
-        buffer = BytesIO(img_bytes)
-        img = Image.open(buffer)
-        img = self.transform(img)
-
-        return img

distributed.py

@@ -1,126 +0,0 @@
-import math
-import pickle
-
-import torch
-from torch import distributed as dist
-from torch.utils.data.sampler import Sampler
-
-
-def get_rank():
-    if not dist.is_available():
-        return 0
-
-    if not dist.is_initialized():
-        return 0
-
-    return dist.get_rank()
-
-
-def synchronize():
-    if not dist.is_available():
-        return
-
-    if not dist.is_initialized():
-        return
-
-    world_size = dist.get_world_size()
-
-    if world_size == 1:
-        return
-
-    dist.barrier()
-
-
-def get_world_size():
-    if not dist.is_available():
-        return 1
-
-    if not dist.is_initialized():
-        return 1
-
-    return dist.get_world_size()
-
-
-def reduce_sum(tensor):
-    if not dist.is_available():
-        return tensor
-
-    if not dist.is_initialized():
-        return tensor
-
-    tensor = tensor.clone()
-    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
-
-    return tensor
-
-
-def gather_grad(params):
-    world_size = get_world_size()
-    
-    if world_size == 1:
-        return
-
-    for param in params:
-        if param.grad is not None:
-            dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
-            param.grad.data.div_(world_size)
-
-
-def all_gather(data):
-    world_size = get_world_size()
-
-    if world_size == 1:
-        return [data]
-
-    buffer = pickle.dumps(data)
-    storage = torch.ByteStorage.from_buffer(buffer)
-    tensor = torch.ByteTensor(storage).to('cuda')
-
-    local_size = torch.IntTensor([tensor.numel()]).to('cuda')
-    size_list = [torch.IntTensor([0]).to('cuda') for _ in range(world_size)]
-    dist.all_gather(size_list, local_size)
-    size_list = [int(size.item()) for size in size_list]
-    max_size = max(size_list)
-
-    tensor_list = []
-    for _ in size_list:
-        tensor_list.append(torch.ByteTensor(size=(max_size,)).to('cuda'))
-
-    if local_size != max_size:
-        padding = torch.ByteTensor(size=(max_size - local_size,)).to('cuda')
-        tensor = torch.cat((tensor, padding), 0)
-
-    dist.all_gather(tensor_list, tensor)
-
-    data_list = []
-
-    for size, tensor in zip(size_list, tensor_list):
-        buffer = tensor.cpu().numpy().tobytes()[:size]
-        data_list.append(pickle.loads(buffer))
-
-    return data_list
-
-
-def reduce_loss_dict(loss_dict):
-    world_size = get_world_size()
-
-    if world_size < 2:
-        return loss_dict
-
-    with torch.no_grad():
-        keys = []
-        losses = []
-
-        for k in sorted(loss_dict.keys()):
-            keys.append(k)
-            losses.append(loss_dict[k])
-
-        losses = torch.stack(losses, 0)
-        dist.reduce(losses, dst=0)
-
-        if dist.get_rank() == 0:
-            losses /= world_size
-
-        reduced_losses = {k: v for k, v in zip(keys, losses)}
-
-    return reduced_losses

download.py

@@ -1,47 +0,0 @@
-from transformers import BertConfig
-
-model_name = 'bert-base-chinese'	# bert版本名称
-model_path = 'D:/Transformers-Bert/bert-base-chinese/'	# 用户下载的预训练bert文件存放地址
-config_path = 'D:/Transformers-Bert/bert-base-chinese/config.json'  # 用户下载的预训练bert文件config.json存放地址
-
-# 载入config 文件可以采取三种方式：bert名称、bert文件夹地址、config文件地址
-# config = BertConfig.from_pretrained(model_name)	# 这个方法会自动从官方的s3数据库下载模型配置、参数等信息（代码中已配置好位置）
-# config = BertConfig.from_pretrained(model_path)
-config = BertConfig.from_pretrained(config_path)
-            
-from transformers import BertTokenizer,BertModel
-model_name = 'KenjieDec/Time-Travel-Rephotograph_e4e_ffhq_encode'
-config = BertConfig.from_pretrained(model_name)	# 这个方法会自动从官方的s3数据库下载模型配置、参数等信息（代码中已配置好位置）
-tokenizer = BertTokenizer.from_pretrained(model_name)		 # 这个方法会自动从官方的s3数据库读取文件下的vocab.txt文件
-model = BertModel.from_pretrained(model_name)		# 这个方法会自动从官方的s3数据库下载模型信息
-
-model_name = 'KenjieDec/Time-Travel-Rephotography_stylegan2-ffhq-config-f'
-config = BertConfig.from_pretrained(model_name)	# 这个方法会自动从官方的s3数据库下载模型配置、参数等信息（代码中已配置好位置）
-tokenizer = BertTokenizer.from_pretrained(model_name)		 # 这个方法会自动从官方的s3数据库读取文件下的vocab.txt文件
-model = BertModel.from_pretrained(model_name)		# 这个方法会自动从官方的s3数据库下载模型信息
-
-model_name = 'KenjieDec/Time-Travel-Rephotography_vgg_face_dag'
-config = BertConfig.from_pretrained(model_name)	# 这个方法会自动从官方的s3数据库下载模型配置、参数等信息（代码中已配置好位置）
-tokenizer = BertTokenizer.from_pretrained(model_name)		 # 这个方法会自动从官方的s3数据库读取文件下的vocab.txt文件
-model = BertModel.from_pretrained(model_name)	
-
-model_name = 'D:\premodel\checkpoint_b'
-config = BertConfig.from_pretrained(model_name)	# 这个方法会自动从官方的s3数据库下载模型配置、参数等信息（代码中已配置好位置）
-tokenizer = BertTokenizer.from_pretrained(model_name)		 # 这个方法会自动从官方的s3数据库读取文件下的vocab.txt文件
-model = BertModel.from_pretrained(model_name)		# 这个方法会自动从官方的s3数据库下载模型信息
-
-model_name = 'D:\premodel\checkpoint_g'
-config = BertConfig.from_pretrained(model_name)	# 这个方法会自动从官方的s3数据库下载模型配置、参数等信息（代码中已配置好位置）
-tokenizer = BertTokenizer.from_pretrained(model_name)		 # 这个方法会自动从官方的s3数据库读取文件下的vocab.txt文件
-model = BertModel.from_pretrained(model_name)		# 这个方法会自动从官方的s3数据库下载模型信息
-
-model_name = 'D:\premodel\checkpoint_gb'
-config = BertConfig.from_pretrained(model_name)	# 这个方法会自动从官方的s3数据库下载模型配置、参数等信息（代码中已配置好位置）
-tokenizer = BertTokenizer.from_pretrained(model_name)		 # 这个方法会自动从官方的s3数据库读取文件下的vocab.txt文件
-model = BertModel.from_pretrained(model_name)		# 这个方法会自动从官方的s3数据库下载模型信息
-
-model_name = 'clearspandex/face-parsing'
-config = BertConfig.from_pretrained(model_name)	# 这个方法会自动从官方的s3数据库下载模型配置、参数等信息（代码中已配置好位置）
-tokenizer = BertTokenizer.from_pretrained(model_name)		 # 这个方法会自动从官方的s3数据库读取文件下的vocab.txt文件
-model = BertModel.from_pretrained(model_name)		# 这个方法会自动从官方的s3数据库下载模型信息
-

fid.py

@@ -1,129 +0,0 @@
-import argparse
-import pickle
-
-import torch
-from torch import nn
-import numpy as np
-from scipy import linalg
-from tqdm import tqdm
-
-from model import Generator
-from calc_inception import load_patched_inception_v3
-
-
-@torch.no_grad()
-def extract_feature_from_samples(
-    generator, inception, truncation, truncation_latent, batch_size, n_sample, device
-):
-    n_batch = n_sample // batch_size
-    resid = n_sample - (n_batch * batch_size)
-    batch_sizes = [batch_size] * n_batch + [resid]
-    features = []
-
-    for batch in tqdm(batch_sizes):
-        latent = torch.randn(batch, 512, device=device)
-        img, _ = g([latent], truncation=truncation, truncation_latent=truncation_latent)
-        feat = inception(img)[0].view(img.shape[0], -1)
-        features.append(feat.to("cpu"))
-
-    features = torch.cat(features, 0)
-
-    return features
-
-
-def calc_fid(sample_mean, sample_cov, real_mean, real_cov, eps=1e-6):
-    cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)
-
-    if not np.isfinite(cov_sqrt).all():
-        print("product of cov matrices is singular")
-        offset = np.eye(sample_cov.shape[0]) * eps
-        cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))
-
-    if np.iscomplexobj(cov_sqrt):
-        if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
-            m = np.max(np.abs(cov_sqrt.imag))
-
-            raise ValueError(f"Imaginary component {m}")
-
-        cov_sqrt = cov_sqrt.real
-
-    mean_diff = sample_mean - real_mean
-    mean_norm = mean_diff @ mean_diff
-
-    trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)
-
-    fid = mean_norm + trace
-
-    return fid
-
-
-if __name__ == "__main__":
-    device = "cuda"
-
-    parser = argparse.ArgumentParser(description="Calculate FID scores")
-
-    parser.add_argument("--truncation", type=float, default=1, help="truncation factor")
-    parser.add_argument(
-        "--truncation_mean",
-        type=int,
-        default=4096,
-        help="number of samples to calculate mean for truncation",
-    )
-    parser.add_argument(
-        "--batch", type=int, default=64, help="batch size for the generator"
-    )
-    parser.add_argument(
-        "--n_sample",
-        type=int,
-        default=50000,
-        help="number of the samples for calculating FID",
-    )
-    parser.add_argument(
-        "--size", type=int, default=256, help="image sizes for generator"
-    )
-    parser.add_argument(
-        "--inception",
-        type=str,
-        default=None,
-        required=True,
-        help="path to precomputed inception embedding",
-    )
-    parser.add_argument(
-        "ckpt", metavar="CHECKPOINT", help="path to generator checkpoint"
-    )
-
-    args = parser.parse_args()
-
-    ckpt = torch.load(args.ckpt)
-
-    g = Generator(args.size, 512, 8).to(device)
-    g.load_state_dict(ckpt["g_ema"])
-    g = nn.DataParallel(g)
-    g.eval()
-
-    if args.truncation < 1:
-        with torch.no_grad():
-            mean_latent = g.mean_latent(args.truncation_mean)
-
-    else:
-        mean_latent = None
-
-    inception = nn.DataParallel(load_patched_inception_v3()).to(device)
-    inception.eval()
-
-    features = extract_feature_from_samples(
-        g, inception, args.truncation, mean_latent, args.batch, args.n_sample, device
-    ).numpy()
-    print(f"extracted {features.shape[0]} features")
-
-    sample_mean = np.mean(features, 0)
-    sample_cov = np.cov(features, rowvar=False)
-
-    with open(args.inception, "rb") as f:
-        embeds = pickle.load(f)
-        real_mean = embeds["mean"]
-        real_cov = embeds["cov"]
-
-    fid = calc_fid(sample_mean, sample_cov, real_mean, real_cov)
-
-    print("fid:", fid)

generate.py

@@ -1,84 +0,0 @@
-import argparse
-
-import torch
-from torchvision import utils
-from model import Generator
-from tqdm import tqdm
-
-
-def generate(args, g_ema, device, mean_latent):
-
-    with torch.no_grad():
-        g_ema.eval()
-        for i in tqdm(range(args.pics)):
-            sample_z = torch.randn(args.sample, args.latent, device=device)
-
-            sample, _ = g_ema(
-                [sample_z], truncation=args.truncation, truncation_latent=mean_latent
-            )
-
-            utils.save_image(
-                sample,
-                f"sample/{str(i).zfill(6)}.png",
-                nrow=1,
-                normalize=True,
-                range=(-1, 1),
-            )
-
-
-if __name__ == "__main__":
-    device = "cuda"
-
-    parser = argparse.ArgumentParser(description="Generate samples from the generator")
-
-    parser.add_argument(
-        "--size", type=int, default=1024, help="output image size of the generator"
-    )
-    parser.add_argument(
-        "--sample",
-        type=int,
-        default=1,
-        help="number of samples to be generated for each image",
-    )
-    parser.add_argument(
-        "--pics", type=int, default=20, help="number of images to be generated"
-    )
-    parser.add_argument("--truncation", type=float, default=1, help="truncation ratio")
-    parser.add_argument(
-        "--truncation_mean",
-        type=int,
-        default=4096,
-        help="number of vectors to calculate mean for the truncation",
-    )
-    parser.add_argument(
-        "--ckpt",
-        type=str,
-        default="stylegan2-ffhq-config-f.pt",
-        help="path to the model checkpoint",
-    )
-    parser.add_argument(
-        "--channel_multiplier",
-        type=int,
-        default=2,
-        help="channel multiplier of the generator. config-f = 2, else = 1",
-    )
-
-    args = parser.parse_args()
-
-    args.latent = 512
-    args.n_mlp = 8
-
-    g_ema = Generator(
-        args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier
-    ).to(device)
-    checkpoint = torch.load(args.ckpt)
-
-    g_ema.load_state_dict(checkpoint["g_ema"])
-
-    if args.truncation < 1:
-        with torch.no_grad():
-            mean_latent = g_ema.mean_latent(args.truncation_mean)
-    else:
-        mean_latent = None
-
-    generate(args, g_ema, device, mean_latent)

inception.py

@@ -1,310 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torchvision import models
-
-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
-
-# 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 InceptionV3(nn.Module):
-    """Pretrained InceptionV3 network returning feature maps"""
-
-    # Index of default block of inception to return,
-    # corresponds to output of final average pooling
-    DEFAULT_BLOCK_INDEX = 3
-
-    # Maps feature dimensionality to their output blocks indices
-    BLOCK_INDEX_BY_DIM = {
-        64: 0,   # First max pooling features
-        192: 1,  # Second max pooling featurs
-        768: 2,  # Pre-aux classifier features
-        2048: 3  # Final average pooling features
-    }
-
-    def __init__(self,
-                 output_blocks=[DEFAULT_BLOCK_INDEX],
-                 resize_input=True,
-                 normalize_input=True,
-                 requires_grad=False,
-                 use_fid_inception=True):
-        """Build pretrained InceptionV3
-
-        Parameters
-        ----------
-        output_blocks : list of int
-            Indices of blocks to return features of. Possible values are:
-                - 0: corresponds to output of first max pooling
-                - 1: corresponds to output of second max pooling
-                - 2: corresponds to output which is fed to aux classifier
-                - 3: corresponds to output of final average pooling
-        resize_input : bool
-            If true, bilinearly resizes input to width and height 299 before
-            feeding input to model. As the network without fully connected
-            layers is fully convolutional, it should be able to handle inputs
-            of arbitrary size, so resizing might not be strictly needed
-        normalize_input : bool
-            If true, scales the input from range (0, 1) to the range the
-            pretrained Inception network expects, namely (-1, 1)
-        requires_grad : bool
-            If true, parameters of the model require gradients. Possibly useful
-            for finetuning the network
-        use_fid_inception : bool
-            If true, uses the pretrained Inception model used in Tensorflow's
-            FID implementation. If false, uses the pretrained Inception model
-            available in torchvision. The FID Inception model has different
-            weights and a slightly different structure from torchvision's
-            Inception model. If you want to compute FID scores, you are
-            strongly advised to set this parameter to true to get comparable
-            results.
-        """
-        super(InceptionV3, self).__init__()
-
-        self.resize_input = resize_input
-        self.normalize_input = normalize_input
-        self.output_blocks = sorted(output_blocks)
-        self.last_needed_block = max(output_blocks)
-
-        assert self.last_needed_block <= 3, \
-            'Last possible output block index is 3'
-
-        self.blocks = nn.ModuleList()
-
-        if use_fid_inception:
-            inception = fid_inception_v3()
-        else:
-            inception = models.inception_v3(pretrained=True)
-
-        # Block 0: input to maxpool1
-        block0 = [
-            inception.Conv2d_1a_3x3,
-            inception.Conv2d_2a_3x3,
-            inception.Conv2d_2b_3x3,
-            nn.MaxPool2d(kernel_size=3, stride=2)
-        ]
-        self.blocks.append(nn.Sequential(*block0))
-
-        # Block 1: maxpool1 to maxpool2
-        if self.last_needed_block >= 1:
-            block1 = [
-                inception.Conv2d_3b_1x1,
-                inception.Conv2d_4a_3x3,
-                nn.MaxPool2d(kernel_size=3, stride=2)
-            ]
-            self.blocks.append(nn.Sequential(*block1))
-
-        # Block 2: maxpool2 to aux classifier
-        if self.last_needed_block >= 2:
-            block2 = [
-                inception.Mixed_5b,
-                inception.Mixed_5c,
-                inception.Mixed_5d,
-                inception.Mixed_6a,
-                inception.Mixed_6b,
-                inception.Mixed_6c,
-                inception.Mixed_6d,
-                inception.Mixed_6e,
-            ]
-            self.blocks.append(nn.Sequential(*block2))
-
-        # Block 3: aux classifier to final avgpool
-        if self.last_needed_block >= 3:
-            block3 = [
-                inception.Mixed_7a,
-                inception.Mixed_7b,
-                inception.Mixed_7c,
-                nn.AdaptiveAvgPool2d(output_size=(1, 1))
-            ]
-            self.blocks.append(nn.Sequential(*block3))
-
-        for param in self.parameters():
-            param.requires_grad = requires_grad
-
-    def forward(self, inp):
-        """Get Inception feature maps
-
-        Parameters
-        ----------
-        inp : torch.autograd.Variable
-            Input tensor of shape Bx3xHxW. Values are expected to be in
-            range (0, 1)
-
-        Returns
-        -------
-        List of torch.autograd.Variable, corresponding to the selected output
-        block, sorted ascending by index
-        """
-        outp = []
-        x = inp
-
-        if self.resize_input:
-            x = F.interpolate(x,
-                              size=(299, 299),
-                              mode='bilinear',
-                              align_corners=False)
-
-        if self.normalize_input:
-            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)
-
-        for idx, block in enumerate(self.blocks):
-            x = block(x)
-            if idx in self.output_blocks:
-                outp.append(x)
-
-            if idx == self.last_needed_block:
-                break
-
-        return outp
-
-
-def fid_inception_v3():
-    """Build pretrained Inception model for FID computation
-
-    The Inception model for FID computation uses a different set of weights
-    and has a slightly different structure than torchvision's Inception.
-
-    This method first constructs torchvision's Inception and then patches the
-    necessary parts that are different in the FID Inception model.
-    """
-    inception = models.inception_v3(num_classes=1008,
-                                    aux_logits=False,
-                                    pretrained=False)
-    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
-    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
-    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
-    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
-    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
-    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
-    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
-    inception.Mixed_7b = FIDInceptionE_1(1280)
-    inception.Mixed_7c = FIDInceptionE_2(2048)
-
-    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
-    inception.load_state_dict(state_dict)
-    return inception
-
-
-class FIDInceptionA(models.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(models.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(models.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(models.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)

lpips/__init__.py

@@ -1,160 +0,0 @@
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-from skimage.measure import compare_ssim
-import torch
-from torch.autograd import Variable
-
-from lpips import dist_model
-
-class PerceptualLoss(torch.nn.Module):
-    def __init__(self, model='net-lin', net='alex', colorspace='rgb', spatial=False, use_gpu=True, gpu_ids=[0]): # VGG using our perceptually-learned weights (LPIPS metric)
-    # def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG as a perceptual loss
-        super(PerceptualLoss, self).__init__()
-        print('Setting up Perceptual loss...')
-        self.use_gpu = use_gpu
-        self.spatial = spatial
-        self.gpu_ids = gpu_ids
-        self.model = dist_model.DistModel()
-        self.model.initialize(model=model, net=net, use_gpu=use_gpu, colorspace=colorspace, spatial=self.spatial, gpu_ids=gpu_ids)
-        print('...[%s] initialized'%self.model.name())
-        print('...Done')
-
-    def forward(self, pred, target, normalize=False):
-        """
-        Pred and target are Variables.
-        If normalize is True, assumes the images are between [0,1] and then scales them between [-1,+1]
-        If normalize is False, assumes the images are already between [-1,+1]
-
-        Inputs pred and target are Nx3xHxW
-        Output pytorch Variable N long
-        """
-
-        if normalize:
-            target = 2 * target  - 1
-            pred = 2 * pred  - 1
-
-        return self.model.forward(target, pred)
-
-def normalize_tensor(in_feat,eps=1e-10):
-    norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1,keepdim=True))
-    return in_feat/(norm_factor+eps)
-
-def l2(p0, p1, range=255.):
-    return .5*np.mean((p0 / range - p1 / range)**2)
-
-def psnr(p0, p1, peak=255.):
-    return 10*np.log10(peak**2/np.mean((1.*p0-1.*p1)**2))
-
-def dssim(p0, p1, range=255.):
-    return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.
-
-def rgb2lab(in_img,mean_cent=False):
-    from skimage import color
-    img_lab = color.rgb2lab(in_img)
-    if(mean_cent):
-        img_lab[:,:,0] = img_lab[:,:,0]-50
-    return img_lab
-
-def tensor2np(tensor_obj):
-    # change dimension of a tensor object into a numpy array
-    return tensor_obj[0].cpu().float().numpy().transpose((1,2,0))
-
-def np2tensor(np_obj):
-     # change dimenion of np array into tensor array
-    return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
-
-def tensor2tensorlab(image_tensor,to_norm=True,mc_only=False):
-    # image tensor to lab tensor
-    from skimage import color
-
-    img = tensor2im(image_tensor)
-    img_lab = color.rgb2lab(img)
-    if(mc_only):
-        img_lab[:,:,0] = img_lab[:,:,0]-50
-    if(to_norm and not mc_only):
-        img_lab[:,:,0] = img_lab[:,:,0]-50
-        img_lab = img_lab/100.
-
-    return np2tensor(img_lab)
-
-def tensorlab2tensor(lab_tensor,return_inbnd=False):
-    from skimage import color
-    import warnings
-    warnings.filterwarnings("ignore")
-
-    lab = tensor2np(lab_tensor)*100.
-    lab[:,:,0] = lab[:,:,0]+50
-
-    rgb_back = 255.*np.clip(color.lab2rgb(lab.astype('float')),0,1)
-    if(return_inbnd):
-        # convert back to lab, see if we match
-        lab_back = color.rgb2lab(rgb_back.astype('uint8'))
-        mask = 1.*np.isclose(lab_back,lab,atol=2.)
-        mask = np2tensor(np.prod(mask,axis=2)[:,:,np.newaxis])
-        return (im2tensor(rgb_back),mask)
-    else:
-        return im2tensor(rgb_back)
-
-def rgb2lab(input):
-    from skimage import color
-    return color.rgb2lab(input / 255.)
-
-def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
-    image_numpy = image_tensor[0].cpu().float().numpy()
-    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
-    return image_numpy.astype(imtype)
-
-def im2tensor(image, imtype=np.uint8, cent=1., factor=255./2.):
-    return torch.Tensor((image / factor - cent)
-                        [:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
-
-def tensor2vec(vector_tensor):
-    return vector_tensor.data.cpu().numpy()[:, :, 0, 0]
-
-def voc_ap(rec, prec, use_07_metric=False):
-    """ ap = voc_ap(rec, prec, [use_07_metric])
-    Compute VOC AP given precision and recall.
-    If use_07_metric is true, uses the
-    VOC 07 11 point method (default:False).
-    """
-    if use_07_metric:
-        # 11 point metric
-        ap = 0.
-        for t in np.arange(0., 1.1, 0.1):
-            if np.sum(rec >= t) == 0:
-                p = 0
-            else:
-                p = np.max(prec[rec >= t])
-            ap = ap + p / 11.
-    else:
-        # correct AP calculation
-        # first append sentinel values at the end
-        mrec = np.concatenate(([0.], rec, [1.]))
-        mpre = np.concatenate(([0.], prec, [0.]))
-
-        # compute the precision envelope
-        for i in range(mpre.size - 1, 0, -1):
-            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
-
-        # to calculate area under PR curve, look for points
-        # where X axis (recall) changes value
-        i = np.where(mrec[1:] != mrec[:-1])[0]
-
-        # and sum (\Delta recall) * prec
-        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
-    return ap
-
-def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
-# def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=1.):
-    image_numpy = image_tensor[0].cpu().float().numpy()
-    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
-    return image_numpy.astype(imtype)
-
-def im2tensor(image, imtype=np.uint8, cent=1., factor=255./2.):
-# def im2tensor(image, imtype=np.uint8, cent=1., factor=1.):
-    return torch.Tensor((image / factor - cent)
-                        [:, :, :, np.newaxis].transpose((3, 2, 0, 1)))

lpips/base_model.py

@@ -1,58 +0,0 @@
-import os
-import numpy as np
-import torch
-from torch.autograd import Variable
-from pdb import set_trace as st
-from IPython import embed
-
-class BaseModel():
-    def __init__(self):
-        pass;
-        
-    def name(self):
-        return 'BaseModel'
-
-    def initialize(self, use_gpu=True, gpu_ids=[0]):
-        self.use_gpu = use_gpu
-        self.gpu_ids = gpu_ids
-
-    def forward(self):
-        pass
-
-    def get_image_paths(self):
-        pass
-
-    def optimize_parameters(self):
-        pass
-
-    def get_current_visuals(self):
-        return self.input
-
-    def get_current_errors(self):
-        return {}
-
-    def save(self, label):
-        pass
-
-    # helper saving function that can be used by subclasses
-    def save_network(self, network, path, network_label, epoch_label):
-        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
-        save_path = os.path.join(path, save_filename)
-        torch.save(network.state_dict(), save_path)
-
-    # helper loading function that can be used by subclasses
-    def load_network(self, network, network_label, epoch_label):
-        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
-        save_path = os.path.join(self.save_dir, save_filename)
-        print('Loading network from %s'%save_path)
-        network.load_state_dict(torch.load(save_path))
-
-    def update_learning_rate():
-        pass
-
-    def get_image_paths(self):
-        return self.image_paths
-
-    def save_done(self, flag=False):
-        np.save(os.path.join(self.save_dir, 'done_flag'),flag)
-        np.savetxt(os.path.join(self.save_dir, 'done_flag'),[flag,],fmt='%i')

lpips/dist_model.py

@@ -1,284 +0,0 @@
-
-from __future__ import absolute_import
-
-import sys
-import numpy as np
-import torch
-from torch import nn
-import os
-from collections import OrderedDict
-from torch.autograd import Variable
-import itertools
-from .base_model import BaseModel
-from scipy.ndimage import zoom
-import fractions
-import functools
-import skimage.transform
-from tqdm import tqdm
-
-from IPython import embed
-
-from . import networks_basic as networks
-import lpips as util
-
-class DistModel(BaseModel):
-    def name(self):
-        return self.model_name
-
-    def initialize(self, model='net-lin', net='alex', colorspace='Lab', pnet_rand=False, pnet_tune=False, model_path=None,
-            use_gpu=True, printNet=False, spatial=False, 
-            is_train=False, lr=.0001, beta1=0.5, version='0.1', gpu_ids=[0]):
-        '''
-        INPUTS
-            model - ['net-lin'] for linearly calibrated network
-                    ['net'] for off-the-shelf network
-                    ['L2'] for L2 distance in Lab colorspace
-                    ['SSIM'] for ssim in RGB colorspace
-            net - ['squeeze','alex','vgg']
-            model_path - if None, will look in weights/[NET_NAME].pth
-            colorspace - ['Lab','RGB'] colorspace to use for L2 and SSIM
-            use_gpu - bool - whether or not to use a GPU
-            printNet - bool - whether or not to print network architecture out
-            spatial - bool - whether to output an array containing varying distances across spatial dimensions
-            spatial_shape - if given, output spatial shape. if None then spatial shape is determined automatically via spatial_factor (see below).
-            spatial_factor - if given, specifies upsampling factor relative to the largest spatial extent of a convolutional layer. if None then resized to size of input images.
-            spatial_order - spline order of filter for upsampling in spatial mode, by default 1 (bilinear).
-            is_train - bool - [True] for training mode
-            lr - float - initial learning rate
-            beta1 - float - initial momentum term for adam
-            version - 0.1 for latest, 0.0 was original (with a bug)
-            gpu_ids - int array - [0] by default, gpus to use
-        '''
-        BaseModel.initialize(self, use_gpu=use_gpu, gpu_ids=gpu_ids)
-
-        self.model = model
-        self.net = net
-        self.is_train = is_train
-        self.spatial = spatial
-        self.gpu_ids = gpu_ids
-        self.model_name = '%s [%s]'%(model,net)
-
-        if(self.model == 'net-lin'): # pretrained net + linear layer
-            self.net = networks.PNetLin(pnet_rand=pnet_rand, pnet_tune=pnet_tune, pnet_type=net,
-                use_dropout=True, spatial=spatial, version=version, lpips=True)
-            kw = {}
-            if not use_gpu:
-                kw['map_location'] = 'cpu'
-            if(model_path is None):
-                import inspect
-                model_path = os.path.abspath(os.path.join(inspect.getfile(self.initialize), '..', 'weights/v%s/%s.pth'%(version,net)))
-
-            if(not is_train):
-                print('Loading model from: %s'%model_path)
-                self.net.load_state_dict(torch.load(model_path, **kw), strict=False)
-
-        elif(self.model=='net'): # pretrained network
-            self.net = networks.PNetLin(pnet_rand=pnet_rand, pnet_type=net, lpips=False)
-        elif(self.model in ['L2','l2']):
-            self.net = networks.L2(use_gpu=use_gpu,colorspace=colorspace) # not really a network, only for testing
-            self.model_name = 'L2'
-        elif(self.model in ['DSSIM','dssim','SSIM','ssim']):
-            self.net = networks.DSSIM(use_gpu=use_gpu,colorspace=colorspace)
-            self.model_name = 'SSIM'
-        else:
-            raise ValueError("Model [%s] not recognized." % self.model)
-
-        self.parameters = list(self.net.parameters())
-
-        if self.is_train: # training mode
-            # extra network on top to go from distances (d0,d1) => predicted human judgment (h*)
-            self.rankLoss = networks.BCERankingLoss()
-            self.parameters += list(self.rankLoss.net.parameters())
-            self.lr = lr
-            self.old_lr = lr
-            self.optimizer_net = torch.optim.Adam(self.parameters, lr=lr, betas=(beta1, 0.999))
-        else: # test mode
-            self.net.eval()
-
-        if(use_gpu):
-            self.net.to(gpu_ids[0])
-            self.net = torch.nn.DataParallel(self.net, device_ids=gpu_ids)
-            if(self.is_train):
-                self.rankLoss = self.rankLoss.to(device=gpu_ids[0]) # just put this on GPU0
-
-        if(printNet):
-            print('---------- Networks initialized -------------')
-            networks.print_network(self.net)
-            print('-----------------------------------------------')
-
-    def forward(self, in0, in1, retPerLayer=False):
-        ''' Function computes the distance between image patches in0 and in1
-        INPUTS
-            in0, in1 - torch.Tensor object of shape Nx3xXxY - image patch scaled to [-1,1]
-        OUTPUT
-            computed distances between in0 and in1
-        '''
-
-        return self.net.forward(in0, in1, retPerLayer=retPerLayer)
-
-    # ***** TRAINING FUNCTIONS *****
-    def optimize_parameters(self):
-        self.forward_train()
-        self.optimizer_net.zero_grad()
-        self.backward_train()
-        self.optimizer_net.step()
-        self.clamp_weights()
-
-    def clamp_weights(self):
-        for module in self.net.modules():
-            if(hasattr(module, 'weight') and module.kernel_size==(1,1)):
-                module.weight.data = torch.clamp(module.weight.data,min=0)
-
-    def set_input(self, data):
-        self.input_ref = data['ref']
-        self.input_p0 = data['p0']
-        self.input_p1 = data['p1']
-        self.input_judge = data['judge']
-
-        if(self.use_gpu):
-            self.input_ref = self.input_ref.to(device=self.gpu_ids[0])
-            self.input_p0 = self.input_p0.to(device=self.gpu_ids[0])
-            self.input_p1 = self.input_p1.to(device=self.gpu_ids[0])
-            self.input_judge = self.input_judge.to(device=self.gpu_ids[0])
-
-        self.var_ref = Variable(self.input_ref,requires_grad=True)
-        self.var_p0 = Variable(self.input_p0,requires_grad=True)
-        self.var_p1 = Variable(self.input_p1,requires_grad=True)
-
-    def forward_train(self): # run forward pass
-        # print(self.net.module.scaling_layer.shift)
-        # print(torch.norm(self.net.module.net.slice1[0].weight).item(), torch.norm(self.net.module.lin0.model[1].weight).item())
-
-        self.d0 = self.forward(self.var_ref, self.var_p0)
-        self.d1 = self.forward(self.var_ref, self.var_p1)
-        self.acc_r = self.compute_accuracy(self.d0,self.d1,self.input_judge)
-
-        self.var_judge = Variable(1.*self.input_judge).view(self.d0.size())
-
-        self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge*2.-1.)
-
-        return self.loss_total
-
-    def backward_train(self):
-        torch.mean(self.loss_total).backward()
-
-    def compute_accuracy(self,d0,d1,judge):
-        ''' d0, d1 are Variables, judge is a Tensor '''
-        d1_lt_d0 = (d1<d0).cpu().data.numpy().flatten()
-        judge_per = judge.cpu().numpy().flatten()
-        return d1_lt_d0*judge_per + (1-d1_lt_d0)*(1-judge_per)
-
-    def get_current_errors(self):
-        retDict = OrderedDict([('loss_total', self.loss_total.data.cpu().numpy()),
-                            ('acc_r', self.acc_r)])
-
-        for key in retDict.keys():
-            retDict[key] = np.mean(retDict[key])
-
-        return retDict
-
-    def get_current_visuals(self):
-        zoom_factor = 256/self.var_ref.data.size()[2]
-
-        ref_img = util.tensor2im(self.var_ref.data)
-        p0_img = util.tensor2im(self.var_p0.data)
-        p1_img = util.tensor2im(self.var_p1.data)
-
-        ref_img_vis = zoom(ref_img,[zoom_factor, zoom_factor, 1],order=0)
-        p0_img_vis = zoom(p0_img,[zoom_factor, zoom_factor, 1],order=0)
-        p1_img_vis = zoom(p1_img,[zoom_factor, zoom_factor, 1],order=0)
-
-        return OrderedDict([('ref', ref_img_vis),
-                            ('p0', p0_img_vis),
-                            ('p1', p1_img_vis)])
-
-    def save(self, path, label):
-        if(self.use_gpu):
-            self.save_network(self.net.module, path, '', label)
-        else:
-            self.save_network(self.net, path, '', label)
-        self.save_network(self.rankLoss.net, path, 'rank', label)
-
-    def update_learning_rate(self,nepoch_decay):
-        lrd = self.lr / nepoch_decay
-        lr = self.old_lr - lrd
-
-        for param_group in self.optimizer_net.param_groups:
-            param_group['lr'] = lr
-
-        print('update lr [%s] decay: %f -> %f' % (type,self.old_lr, lr))
-        self.old_lr = lr
-
-def score_2afc_dataset(data_loader, func, name=''):
-    ''' Function computes Two Alternative Forced Choice (2AFC) score using
-        distance function 'func' in dataset 'data_loader'
-    INPUTS
-        data_loader - CustomDatasetDataLoader object - contains a TwoAFCDataset inside
-        func - callable distance function - calling d=func(in0,in1) should take 2
-            pytorch tensors with shape Nx3xXxY, and return numpy array of length N
-    OUTPUTS
-        [0] - 2AFC score in [0,1], fraction of time func agrees with human evaluators
-        [1] - dictionary with following elements
-            d0s,d1s - N arrays containing distances between reference patch to perturbed patches 
-            gts - N array in [0,1], preferred patch selected by human evaluators
-                (closer to "0" for left patch p0, "1" for right patch p1,
-                "0.6" means 60pct people preferred right patch, 40pct preferred left)
-            scores - N array in [0,1], corresponding to what percentage function agreed with humans
-    CONSTS
-        N - number of test triplets in data_loader
-    '''
-
-    d0s = []
-    d1s = []
-    gts = []
-
-    for data in tqdm(data_loader.load_data(), desc=name):
-        d0s+=func(data['ref'],data['p0']).data.cpu().numpy().flatten().tolist()
-        d1s+=func(data['ref'],data['p1']).data.cpu().numpy().flatten().tolist()
-        gts+=data['judge'].cpu().numpy().flatten().tolist()
-
-    d0s = np.array(d0s)
-    d1s = np.array(d1s)
-    gts = np.array(gts)
-    scores = (d0s<d1s)*(1.-gts) + (d1s<d0s)*gts + (d1s==d0s)*.5
-
-    return(np.mean(scores), dict(d0s=d0s,d1s=d1s,gts=gts,scores=scores))
-
-def score_jnd_dataset(data_loader, func, name=''):
-    ''' Function computes JND score using distance function 'func' in dataset 'data_loader'
-    INPUTS
-        data_loader - CustomDatasetDataLoader object - contains a JNDDataset inside
-        func - callable distance function - calling d=func(in0,in1) should take 2
-            pytorch tensors with shape Nx3xXxY, and return pytorch array of length N
-    OUTPUTS
-        [0] - JND score in [0,1], mAP score (area under precision-recall curve)
-        [1] - dictionary with following elements
-            ds - N array containing distances between two patches shown to human evaluator
-            sames - N array containing fraction of people who thought the two patches were identical
-    CONSTS
-        N - number of test triplets in data_loader
-    '''
-
-    ds = []
-    gts = []
-
-    for data in tqdm(data_loader.load_data(), desc=name):
-        ds+=func(data['p0'],data['p1']).data.cpu().numpy().tolist()
-        gts+=data['same'].cpu().numpy().flatten().tolist()
-
-    sames = np.array(gts)
-    ds = np.array(ds)
-
-    sorted_inds = np.argsort(ds)
-    ds_sorted = ds[sorted_inds]
-    sames_sorted = sames[sorted_inds]
-
-    TPs = np.cumsum(sames_sorted)
-    FPs = np.cumsum(1-sames_sorted)
-    FNs = np.sum(sames_sorted)-TPs
-
-    precs = TPs/(TPs+FPs)
-    recs = TPs/(TPs+FNs)
-    score = util.voc_ap(recs,precs)
-
-    return(score, dict(ds=ds,sames=sames))

lpips/networks_basic.py

@@ -1,187 +0,0 @@
-
-from __future__ import absolute_import
-
-import sys
-import torch
-import torch.nn as nn
-import torch.nn.init as init
-from torch.autograd import Variable
-import numpy as np
-from pdb import set_trace as st
-from skimage import color
-from IPython import embed
-from . import pretrained_networks as pn
-
-import lpips as util
-
-def spatial_average(in_tens, keepdim=True):
-    return in_tens.mean([2,3],keepdim=keepdim)
-
-def upsample(in_tens, out_H=64): # assumes scale factor is same for H and W
-    in_H = in_tens.shape[2]
-    scale_factor = 1.*out_H/in_H
-
-    return nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=False)(in_tens)
-
-# Learned perceptual metric
-class PNetLin(nn.Module):
-    def __init__(self, pnet_type='vgg', pnet_rand=False, pnet_tune=False, use_dropout=True, spatial=False, version='0.1', lpips=True):
-        super(PNetLin, self).__init__()
-
-        self.pnet_type = pnet_type
-        self.pnet_tune = pnet_tune
-        self.pnet_rand = pnet_rand
-        self.spatial = spatial
-        self.lpips = lpips
-        self.version = version
-        self.scaling_layer = ScalingLayer()
-
-        if(self.pnet_type in ['vgg','vgg16']):
-            net_type = pn.vgg16
-            self.chns = [64,128,256,512,512]
-        elif(self.pnet_type=='alex'):
-            net_type = pn.alexnet
-            self.chns = [64,192,384,256,256]
-        elif(self.pnet_type=='squeeze'):
-            net_type = pn.squeezenet
-            self.chns = [64,128,256,384,384,512,512]
-        self.L = len(self.chns)
-
-        self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune)
-
-        if(lpips):
-            self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
-            self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
-            self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
-            self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
-            self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
-            self.lins = [self.lin0,self.lin1,self.lin2,self.lin3,self.lin4]
-            if(self.pnet_type=='squeeze'): # 7 layers for squeezenet
-                self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout)
-                self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout)
-                self.lins+=[self.lin5,self.lin6]
-
-    def forward(self, in0, in1, retPerLayer=False):
-        # v0.0 - original release had a bug, where input was not scaled
-        in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1)
-        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
-        feats0, feats1, diffs = {}, {}, {}
-
-        for kk in range(self.L):
-            feats0[kk], feats1[kk] = util.normalize_tensor(outs0[kk]), util.normalize_tensor(outs1[kk])
-            diffs[kk] = (feats0[kk]-feats1[kk])**2
-
-        if(self.lpips):
-            if(self.spatial):
-                res = [upsample(self.lins[kk].model(diffs[kk]), out_H=in0.shape[2]) for kk in range(self.L)]
-            else:
-                res = [spatial_average(self.lins[kk].model(diffs[kk]), keepdim=True) for kk in range(self.L)]
-        else:
-            if(self.spatial):
-                res = [upsample(diffs[kk].sum(dim=1,keepdim=True), out_H=in0.shape[2]) for kk in range(self.L)]
-            else:
-                res = [spatial_average(diffs[kk].sum(dim=1,keepdim=True), keepdim=True) for kk in range(self.L)]
-
-        val = res[0]
-        for l in range(1,self.L):
-            val += res[l]
-        
-        if(retPerLayer):
-            return (val, res)
-        else:
-            return val
-
-class ScalingLayer(nn.Module):
-    def __init__(self):
-        super(ScalingLayer, self).__init__()
-        self.register_buffer('shift', torch.Tensor([-.030,-.088,-.188])[None,:,None,None])
-        self.register_buffer('scale', torch.Tensor([.458,.448,.450])[None,:,None,None])
-
-    def forward(self, inp):
-        return (inp - self.shift) / self.scale
-
-
-class NetLinLayer(nn.Module):
-    ''' A single linear layer which does a 1x1 conv '''
-    def __init__(self, chn_in, chn_out=1, use_dropout=False):
-        super(NetLinLayer, self).__init__()
-
-        layers = [nn.Dropout(),] if(use_dropout) else []
-        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),]
-        self.model = nn.Sequential(*layers)
-
-
-class Dist2LogitLayer(nn.Module):
-    ''' takes 2 distances, puts through fc layers, spits out value between [0,1] (if use_sigmoid is True) '''
-    def __init__(self, chn_mid=32, use_sigmoid=True):
-        super(Dist2LogitLayer, self).__init__()
-
-        layers = [nn.Conv2d(5, chn_mid, 1, stride=1, padding=0, bias=True),]
-        layers += [nn.LeakyReLU(0.2,True),]
-        layers += [nn.Conv2d(chn_mid, chn_mid, 1, stride=1, padding=0, bias=True),]
-        layers += [nn.LeakyReLU(0.2,True),]
-        layers += [nn.Conv2d(chn_mid, 1, 1, stride=1, padding=0, bias=True),]
-        if(use_sigmoid):
-            layers += [nn.Sigmoid(),]
-        self.model = nn.Sequential(*layers)
-
-    def forward(self,d0,d1,eps=0.1):
-        return self.model.forward(torch.cat((d0,d1,d0-d1,d0/(d1+eps),d1/(d0+eps)),dim=1))
-
-class BCERankingLoss(nn.Module):
-    def __init__(self, chn_mid=32):
-        super(BCERankingLoss, self).__init__()
-        self.net = Dist2LogitLayer(chn_mid=chn_mid)
-        # self.parameters = list(self.net.parameters())
-        self.loss = torch.nn.BCELoss()
-
-    def forward(self, d0, d1, judge):
-        per = (judge+1.)/2.
-        self.logit = self.net.forward(d0,d1)
-        return self.loss(self.logit, per)
-
-# L2, DSSIM metrics
-class FakeNet(nn.Module):
-    def __init__(self, use_gpu=True, colorspace='Lab'):
-        super(FakeNet, self).__init__()
-        self.use_gpu = use_gpu
-        self.colorspace=colorspace
-
-class L2(FakeNet):
-
-    def forward(self, in0, in1, retPerLayer=None):
-        assert(in0.size()[0]==1) # currently only supports batchSize 1
-
-        if(self.colorspace=='RGB'):
-            (N,C,X,Y) = in0.size()
-            value = torch.mean(torch.mean(torch.mean((in0-in1)**2,dim=1).view(N,1,X,Y),dim=2).view(N,1,1,Y),dim=3).view(N)
-            return value
-        elif(self.colorspace=='Lab'):
-            value = util.l2(util.tensor2np(util.tensor2tensorlab(in0.data,to_norm=False)), 
-                util.tensor2np(util.tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
-            ret_var = Variable( torch.Tensor((value,) ) )
-            if(self.use_gpu):
-                ret_var = ret_var.cuda()
-            return ret_var
-
-class DSSIM(FakeNet):
-
-    def forward(self, in0, in1, retPerLayer=None):
-        assert(in0.size()[0]==1) # currently only supports batchSize 1
-
-        if(self.colorspace=='RGB'):
-            value = util.dssim(1.*util.tensor2im(in0.data), 1.*util.tensor2im(in1.data), range=255.).astype('float')
-        elif(self.colorspace=='Lab'):
-            value = util.dssim(util.tensor2np(util.tensor2tensorlab(in0.data,to_norm=False)), 
-                util.tensor2np(util.tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
-        ret_var = Variable( torch.Tensor((value,) ) )
-        if(self.use_gpu):
-            ret_var = ret_var.cuda()
-        return ret_var
-
-def print_network(net):
-    num_params = 0
-    for param in net.parameters():
-        num_params += param.numel()
-    print('Network',net)
-    print('Total number of parameters: %d' % num_params)

lpips/pretrained_networks.py

@@ -1,181 +0,0 @@
-from collections import namedtuple
-import torch
-from torchvision import models as tv
-from IPython import embed
-
-class squeezenet(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True):
-        super(squeezenet, self).__init__()
-        pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        self.slice6 = torch.nn.Sequential()
-        self.slice7 = torch.nn.Sequential()
-        self.N_slices = 7
-        for x in range(2):
-            self.slice1.add_module(str(x), pretrained_features[x])
-        for x in range(2,5):
-            self.slice2.add_module(str(x), pretrained_features[x])
-        for x in range(5, 8):
-            self.slice3.add_module(str(x), pretrained_features[x])
-        for x in range(8, 10):
-            self.slice4.add_module(str(x), pretrained_features[x])
-        for x in range(10, 11):
-            self.slice5.add_module(str(x), pretrained_features[x])
-        for x in range(11, 12):
-            self.slice6.add_module(str(x), pretrained_features[x])
-        for x in range(12, 13):
-            self.slice7.add_module(str(x), pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1 = h
-        h = self.slice2(h)
-        h_relu2 = h
-        h = self.slice3(h)
-        h_relu3 = h
-        h = self.slice4(h)
-        h_relu4 = h
-        h = self.slice5(h)
-        h_relu5 = h
-        h = self.slice6(h)
-        h_relu6 = h
-        h = self.slice7(h)
-        h_relu7 = h
-        vgg_outputs = namedtuple("SqueezeOutputs", ['relu1','relu2','relu3','relu4','relu5','relu6','relu7'])
-        out = vgg_outputs(h_relu1,h_relu2,h_relu3,h_relu4,h_relu5,h_relu6,h_relu7)
-
-        return out
-
-
-class alexnet(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True):
-        super(alexnet, self).__init__()
-        alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        self.N_slices = 5
-        for x in range(2):
-            self.slice1.add_module(str(x), alexnet_pretrained_features[x])
-        for x in range(2, 5):
-            self.slice2.add_module(str(x), alexnet_pretrained_features[x])
-        for x in range(5, 8):
-            self.slice3.add_module(str(x), alexnet_pretrained_features[x])
-        for x in range(8, 10):
-            self.slice4.add_module(str(x), alexnet_pretrained_features[x])
-        for x in range(10, 12):
-            self.slice5.add_module(str(x), alexnet_pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1 = h
-        h = self.slice2(h)
-        h_relu2 = h
-        h = self.slice3(h)
-        h_relu3 = h
-        h = self.slice4(h)
-        h_relu4 = h
-        h = self.slice5(h)
-        h_relu5 = h
-        alexnet_outputs = namedtuple("AlexnetOutputs", ['relu1', 'relu2', 'relu3', 'relu4', 'relu5'])
-        out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)
-
-        return out
-
-class vgg16(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True):
-        super(vgg16, self).__init__()
-        vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        self.N_slices = 5
-        for x in range(4):
-            self.slice1.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(4, 9):
-            self.slice2.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(9, 16):
-            self.slice3.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(16, 23):
-            self.slice4.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(23, 30):
-            self.slice5.add_module(str(x), vgg_pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1_2 = h
-        h = self.slice2(h)
-        h_relu2_2 = h
-        h = self.slice3(h)
-        h_relu3_3 = h
-        h = self.slice4(h)
-        h_relu4_3 = h
-        h = self.slice5(h)
-        h_relu5_3 = h
-        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
-        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
-
-        return out
-
-
-
-class resnet(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True, num=18):
-        super(resnet, self).__init__()
-        if(num==18):
-            self.net = tv.resnet18(pretrained=pretrained)
-        elif(num==34):
-            self.net = tv.resnet34(pretrained=pretrained)
-        elif(num==50):
-            self.net = tv.resnet50(pretrained=pretrained)
-        elif(num==101):
-            self.net = tv.resnet101(pretrained=pretrained)
-        elif(num==152):
-            self.net = tv.resnet152(pretrained=pretrained)
-        self.N_slices = 5
-
-        self.conv1 = self.net.conv1
-        self.bn1 = self.net.bn1
-        self.relu = self.net.relu
-        self.maxpool = self.net.maxpool
-        self.layer1 = self.net.layer1
-        self.layer2 = self.net.layer2
-        self.layer3 = self.net.layer3
-        self.layer4 = self.net.layer4
-
-    def forward(self, X):
-        h = self.conv1(X)
-        h = self.bn1(h)
-        h = self.relu(h)
-        h_relu1 = h
-        h = self.maxpool(h)
-        h = self.layer1(h)
-        h_conv2 = h
-        h = self.layer2(h)
-        h_conv3 = h
-        h = self.layer3(h)
-        h_conv4 = h
-        h = self.layer4(h)
-        h_conv5 = h
-
-        outputs = namedtuple("Outputs", ['relu1','conv2','conv3','conv4','conv5'])
-        out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)
-
-        return out

lpips/weights/v0.0/alex.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:18720f55913d0af89042f13faa7e536a6ce1444a0914e6db9461355ece1e8cd5
-size 5455

lpips/weights/v0.0/squeeze.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c27abd3a0145541baa50990817df58d3759c3f8154949f42af3b59b4e042d0bf
-size 10057

lpips/weights/v0.0/vgg.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b9e4236260c3dd988fc79d2a48d645d885afcbb21f9fd595e6744cf7419b582c
-size 6735

lpips/weights/v0.1/alex.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:df73285e35b22355a2df87cdb6b70b343713b667eddbda73e1977e0c860835c0
-size 6009

lpips/weights/v0.1/squeeze.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4a5350f23600cb79923ce65bb07cbf57dca461329894153e05a1346bd531cf76
-size 10811

lpips/weights/v0.1/vgg.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a78928a0af1e5f0fcb1f3b9e8f8c3a2a5a3de244d830ad5c1feddc79b8432868
-size 7289

model.py

@@ -1,698 +0,0 @@
-import math
-import random
-import functools
-import operator
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.autograd import Function
-
-from op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d, conv2d_gradfix
-
-
-class PixelNorm(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, input):
-        return input * torch.rsqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-8)
-
-
-def make_kernel(k):
-    k = torch.tensor(k, dtype=torch.float32)
-
-    if k.ndim == 1:
-        k = k[None, :] * k[:, None]
-
-    k /= k.sum()
-
-    return k
-
-
-class Upsample(nn.Module):
-    def __init__(self, kernel, factor=2):
-        super().__init__()
-
-        self.factor = factor
-        kernel = make_kernel(kernel) * (factor ** 2)
-        self.register_buffer("kernel", kernel)
-
-        p = kernel.shape[0] - factor
-
-        pad0 = (p + 1) // 2 + factor - 1
-        pad1 = p // 2
-
-        self.pad = (pad0, pad1)
-
-    def forward(self, input):
-        out = upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad)
-
-        return out
-
-
-class Downsample(nn.Module):
-    def __init__(self, kernel, factor=2):
-        super().__init__()
-
-        self.factor = factor
-        kernel = make_kernel(kernel)
-        self.register_buffer("kernel", kernel)
-
-        p = kernel.shape[0] - factor
-
-        pad0 = (p + 1) // 2
-        pad1 = p // 2
-
-        self.pad = (pad0, pad1)
-
-    def forward(self, input):
-        out = upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)
-
-        return out
-
-
-class Blur(nn.Module):
-    def __init__(self, kernel, pad, upsample_factor=1):
-        super().__init__()
-
-        kernel = make_kernel(kernel)
-
-        if upsample_factor > 1:
-            kernel = kernel * (upsample_factor ** 2)
-
-        self.register_buffer("kernel", kernel)
-
-        self.pad = pad
-
-    def forward(self, input):
-        out = upfirdn2d(input, self.kernel, pad=self.pad)
-
-        return out
-
-
-class EqualConv2d(nn.Module):
-    def __init__(
-        self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True
-    ):
-        super().__init__()
-
-        self.weight = nn.Parameter(
-            torch.randn(out_channel, in_channel, kernel_size, kernel_size)
-        )
-        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
-
-        self.stride = stride
-        self.padding = padding
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(out_channel))
-
-        else:
-            self.bias = None
-
-    def forward(self, input):
-        out = conv2d_gradfix.conv2d(
-            input,
-            self.weight * self.scale,
-            bias=self.bias,
-            stride=self.stride,
-            padding=self.padding,
-        )
-
-        return out
-
-    def __repr__(self):
-        return (
-            f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},"
-            f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})"
-        )
-
-
-class EqualLinear(nn.Module):
-    def __init__(
-        self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None
-    ):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
-
-        else:
-            self.bias = None
-
-        self.activation = activation
-
-        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
-        self.lr_mul = lr_mul
-
-    def forward(self, input):
-        if self.activation:
-            out = F.linear(input, self.weight * self.scale)
-            out = fused_leaky_relu(out, self.bias * self.lr_mul)
-
-        else:
-            out = F.linear(
-                input, self.weight * self.scale, bias=self.bias * self.lr_mul
-            )
-
-        return out
-
-    def __repr__(self):
-        return (
-            f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})"
-        )
-
-
-class ModulatedConv2d(nn.Module):
-    def __init__(
-        self,
-        in_channel,
-        out_channel,
-        kernel_size,
-        style_dim,
-        demodulate=True,
-        upsample=False,
-        downsample=False,
-        blur_kernel=[1, 3, 3, 1],
-        fused=True,
-    ):
-        super().__init__()
-
-        self.eps = 1e-8
-        self.kernel_size = kernel_size
-        self.in_channel = in_channel
-        self.out_channel = out_channel
-        self.upsample = upsample
-        self.downsample = downsample
-
-        if upsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) - (kernel_size - 1)
-            pad0 = (p + 1) // 2 + factor - 1
-            pad1 = p // 2 + 1
-
-            self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)
-
-        if downsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) + (kernel_size - 1)
-            pad0 = (p + 1) // 2
-            pad1 = p // 2
-
-            self.blur = Blur(blur_kernel, pad=(pad0, pad1))
-
-        fan_in = in_channel * kernel_size ** 2
-        self.scale = 1 / math.sqrt(fan_in)
-        self.padding = kernel_size // 2
-
-        self.weight = nn.Parameter(
-            torch.randn(1, out_channel, in_channel, kernel_size, kernel_size)
-        )
-
-        self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)
-
-        self.demodulate = demodulate
-        self.fused = fused
-
-    def __repr__(self):
-        return (
-            f"{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, "
-            f"upsample={self.upsample}, downsample={self.downsample})"
-        )
-
-    def forward(self, input, style):
-        batch, in_channel, height, width = input.shape
-
-        if not self.fused:
-            weight = self.scale * self.weight.squeeze(0)
-            style = self.modulation(style)
-
-            if self.demodulate:
-                w = weight.unsqueeze(0) * style.view(batch, 1, in_channel, 1, 1)
-                dcoefs = (w.square().sum((2, 3, 4)) + 1e-8).rsqrt()
-
-            input = input * style.reshape(batch, in_channel, 1, 1)
-
-            if self.upsample:
-                weight = weight.transpose(0, 1)
-                out = conv2d_gradfix.conv_transpose2d(
-                    input, weight, padding=0, stride=2
-                )
-                out = self.blur(out)
-
-            elif self.downsample:
-                input = self.blur(input)
-                out = conv2d_gradfix.conv2d(input, weight, padding=0, stride=2)
-
-            else:
-                out = conv2d_gradfix.conv2d(input, weight, padding=self.padding)
-
-            if self.demodulate:
-                out = out * dcoefs.view(batch, -1, 1, 1)
-
-            return out
-
-        style = self.modulation(style).view(batch, 1, in_channel, 1, 1)
-        weight = self.scale * self.weight * style
-
-        if self.demodulate:
-            demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
-            weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)
-
-        weight = weight.view(
-            batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size
-        )
-
-        if self.upsample:
-            input = input.view(1, batch * in_channel, height, width)
-            weight = weight.view(
-                batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size
-            )
-            weight = weight.transpose(1, 2).reshape(
-                batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size
-            )
-            out = conv2d_gradfix.conv_transpose2d(
-                input, weight, padding=0, stride=2, groups=batch
-            )
-            _, _, height, width = out.shape
-            out = out.view(batch, self.out_channel, height, width)
-            out = self.blur(out)
-
-        elif self.downsample:
-            input = self.blur(input)
-            _, _, height, width = input.shape
-            input = input.view(1, batch * in_channel, height, width)
-            out = conv2d_gradfix.conv2d(
-                input, weight, padding=0, stride=2, groups=batch
-            )
-            _, _, height, width = out.shape
-            out = out.view(batch, self.out_channel, height, width)
-
-        else:
-            input = input.view(1, batch * in_channel, height, width)
-            out = conv2d_gradfix.conv2d(
-                input, weight, padding=self.padding, groups=batch
-            )
-            _, _, height, width = out.shape
-            out = out.view(batch, self.out_channel, height, width)
-
-        return out
-
-
-class NoiseInjection(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        self.weight = nn.Parameter(torch.zeros(1))
-
-    def forward(self, image, noise=None):
-        if noise is None:
-            batch, _, height, width = image.shape
-            noise = image.new_empty(batch, 1, height, width).normal_()
-
-        return image + self.weight * noise
-
-
-class ConstantInput(nn.Module):
-    def __init__(self, channel, size=4):
-        super().__init__()
-
-        self.input = nn.Parameter(torch.randn(1, channel, size, size))
-
-    def forward(self, input):
-        batch = input.shape[0]
-        out = self.input.repeat(batch, 1, 1, 1)
-
-        return out
-
-
-class StyledConv(nn.Module):
-    def __init__(
-        self,
-        in_channel,
-        out_channel,
-        kernel_size,
-        style_dim,
-        upsample=False,
-        blur_kernel=[1, 3, 3, 1],
-        demodulate=True,
-    ):
-        super().__init__()
-
-        self.conv = ModulatedConv2d(
-            in_channel,
-            out_channel,
-            kernel_size,
-            style_dim,
-            upsample=upsample,
-            blur_kernel=blur_kernel,
-            demodulate=demodulate,
-        )
-
-        self.noise = NoiseInjection()
-        # self.bias = nn.Parameter(torch.zeros(1, out_channel, 1, 1))
-        # self.activate = ScaledLeakyReLU(0.2)
-        self.activate = FusedLeakyReLU(out_channel)
-
-    def forward(self, input, style, noise=None):
-        out = self.conv(input, style)
-        out = self.noise(out, noise=noise)
-        # out = out + self.bias
-        out = self.activate(out)
-
-        return out
-
-
-class ToRGB(nn.Module):
-    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        if upsample:
-            self.upsample = Upsample(blur_kernel)
-
-        self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
-        self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
-
-    def forward(self, input, style, skip=None):
-        out = self.conv(input, style)
-        out = out + self.bias
-
-        if skip is not None:
-            skip = self.upsample(skip)
-
-            out = out + skip
-
-        return out
-
-
-class Generator(nn.Module):
-    def __init__(
-        self,
-        size,
-        style_dim,
-        n_mlp,
-        channel_multiplier=2,
-        blur_kernel=[1, 3, 3, 1],
-        lr_mlp=0.01,
-    ):
-        super().__init__()
-
-        self.size = size
-
-        self.style_dim = style_dim
-
-        layers = [PixelNorm()]
-
-        for i in range(n_mlp):
-            layers.append(
-                EqualLinear(
-                    style_dim, style_dim, lr_mul=lr_mlp, activation="fused_lrelu"
-                )
-            )
-
-        self.style = nn.Sequential(*layers)
-
-        self.channels = {
-            4: 512,
-            8: 512,
-            16: 512,
-            32: 512,
-            64: 256 * channel_multiplier,
-            128: 128 * channel_multiplier,
-            256: 64 * channel_multiplier,
-            512: 32 * channel_multiplier,
-            1024: 16 * channel_multiplier,
-        }
-
-        self.input = ConstantInput(self.channels[4])
-        self.conv1 = StyledConv(
-            self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel
-        )
-        self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)
-
-        self.log_size = int(math.log(size, 2))
-        self.num_layers = (self.log_size - 2) * 2 + 1
-
-        self.convs = nn.ModuleList()
-        self.upsamples = nn.ModuleList()
-        self.to_rgbs = nn.ModuleList()
-        self.noises = nn.Module()
-
-        in_channel = self.channels[4]
-
-        for layer_idx in range(self.num_layers):
-            res = (layer_idx + 5) // 2
-            shape = [1, 1, 2 ** res, 2 ** res]
-            self.noises.register_buffer(f"noise_{layer_idx}", torch.randn(*shape))
-
-        for i in range(3, self.log_size + 1):
-            out_channel = self.channels[2 ** i]
-
-            self.convs.append(
-                StyledConv(
-                    in_channel,
-                    out_channel,
-                    3,
-                    style_dim,
-                    upsample=True,
-                    blur_kernel=blur_kernel,
-                )
-            )
-
-            self.convs.append(
-                StyledConv(
-                    out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel
-                )
-            )
-
-            self.to_rgbs.append(ToRGB(out_channel, style_dim))
-
-            in_channel = out_channel
-
-        self.n_latent = self.log_size * 2 - 2
-
-    def make_noise(self):
-        device = self.input.input.device
-
-        noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]
-
-        for i in range(3, self.log_size + 1):
-            for _ in range(2):
-                noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device))
-
-        return noises
-
-    def mean_latent(self, n_latent):
-        latent_in = torch.randn(
-            n_latent, self.style_dim, device=self.input.input.device
-        )
-        latent = self.style(latent_in).mean(0, keepdim=True)
-
-        return latent
-
-    def get_latent(self, input):
-        return self.style(input)
-
-    def forward(
-        self,
-        styles,
-        return_latents=False,
-        inject_index=None,
-        truncation=1,
-        truncation_latent=None,
-        input_is_latent=False,
-        noise=None,
-        randomize_noise=True,
-    ):
-        if not input_is_latent:
-            styles = [self.style(s) for s in styles]
-
-        if noise is None:
-            if randomize_noise:
-                noise = [None] * self.num_layers
-            else:
-                noise = [
-                    getattr(self.noises, f"noise_{i}") for i in range(self.num_layers)
-                ]
-
-        if truncation < 1:
-            style_t = []
-
-            for style in styles:
-                style_t.append(
-                    truncation_latent + truncation * (style - truncation_latent)
-                )
-
-            styles = style_t
-
-        if len(styles) < 2:
-            inject_index = self.n_latent
-
-            if styles[0].ndim < 3:
-                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
-
-            else:
-                latent = styles[0]
-
-        else:
-            if inject_index is None:
-                inject_index = random.randint(1, self.n_latent - 1)
-
-            latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
-            latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)
-
-            latent = torch.cat([latent, latent2], 1)
-
-        out = self.input(latent)
-        out = self.conv1(out, latent[:, 0], noise=noise[0])
-
-        skip = self.to_rgb1(out, latent[:, 1])
-
-        i = 1
-        for conv1, conv2, noise1, noise2, to_rgb in zip(
-            self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
-        ):
-            out = conv1(out, latent[:, i], noise=noise1)
-            out = conv2(out, latent[:, i + 1], noise=noise2)
-            skip = to_rgb(out, latent[:, i + 2], skip)
-
-            i += 2
-
-        image = skip
-
-        if return_latents:
-            return image, latent
-
-        else:
-            return image, None
-
-
-class ConvLayer(nn.Sequential):
-    def __init__(
-        self,
-        in_channel,
-        out_channel,
-        kernel_size,
-        downsample=False,
-        blur_kernel=[1, 3, 3, 1],
-        bias=True,
-        activate=True,
-    ):
-        layers = []
-
-        if downsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) + (kernel_size - 1)
-            pad0 = (p + 1) // 2
-            pad1 = p // 2
-
-            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
-
-            stride = 2
-            self.padding = 0
-
-        else:
-            stride = 1
-            self.padding = kernel_size // 2
-
-        layers.append(
-            EqualConv2d(
-                in_channel,
-                out_channel,
-                kernel_size,
-                padding=self.padding,
-                stride=stride,
-                bias=bias and not activate,
-            )
-        )
-
-        if activate:
-            layers.append(FusedLeakyReLU(out_channel, bias=bias))
-
-        super().__init__(*layers)
-
-
-class ResBlock(nn.Module):
-    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        self.conv1 = ConvLayer(in_channel, in_channel, 3)
-        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
-
-        self.skip = ConvLayer(
-            in_channel, out_channel, 1, downsample=True, activate=False, bias=False
-        )
-
-    def forward(self, input):
-        out = self.conv1(input)
-        out = self.conv2(out)
-
-        skip = self.skip(input)
-        out = (out + skip) / math.sqrt(2)
-
-        return out
-
-
-class Discriminator(nn.Module):
-    def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        channels = {
-            4: 512,
-            8: 512,
-            16: 512,
-            32: 512,
-            64: 256 * channel_multiplier,
-            128: 128 * channel_multiplier,
-            256: 64 * channel_multiplier,
-            512: 32 * channel_multiplier,
-            1024: 16 * channel_multiplier,
-        }
-
-        convs = [ConvLayer(3, channels[size], 1)]
-
-        log_size = int(math.log(size, 2))
-
-        in_channel = channels[size]
-
-        for i in range(log_size, 2, -1):
-            out_channel = channels[2 ** (i - 1)]
-
-            convs.append(ResBlock(in_channel, out_channel, blur_kernel))
-
-            in_channel = out_channel
-
-        self.convs = nn.Sequential(*convs)
-
-        self.stddev_group = 4
-        self.stddev_feat = 1
-
-        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
-        self.final_linear = nn.Sequential(
-            EqualLinear(channels[4] * 4 * 4, channels[4], activation="fused_lrelu"),
-            EqualLinear(channels[4], 1),
-        )
-
-    def forward(self, input):
-        out = self.convs(input)
-
-        batch, channel, height, width = out.shape
-        group = min(batch, self.stddev_group)
-        stddev = out.view(
-            group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
-        )
-        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
-        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
-        stddev = stddev.repeat(group, 1, height, width)
-        out = torch.cat([out, stddev], 1)
-
-        out = self.final_conv(out)
-
-        out = out.view(batch, -1)
-        out = self.final_linear(out)
-
-        return out
-

non_leaking.py

@@ -1,465 +0,0 @@
-import math
-
-import torch
-from torch import autograd
-from torch.nn import functional as F
-import numpy as np
-
-from distributed import reduce_sum
-from op import upfirdn2d
-
-
-class AdaptiveAugment:
-    def __init__(self, ada_aug_target, ada_aug_len, update_every, device):
-        self.ada_aug_target = ada_aug_target
-        self.ada_aug_len = ada_aug_len
-        self.update_every = update_every
-
-        self.ada_update = 0
-        self.ada_aug_buf = torch.tensor([0.0, 0.0], device=device)
-        self.r_t_stat = 0
-        self.ada_aug_p = 0
-
-    @torch.no_grad()
-    def tune(self, real_pred):
-        self.ada_aug_buf += torch.tensor(
-            (torch.sign(real_pred).sum().item(), real_pred.shape[0]),
-            device=real_pred.device,
-        )
-        self.ada_update += 1
-
-        if self.ada_update % self.update_every == 0:
-            self.ada_aug_buf = reduce_sum(self.ada_aug_buf)
-            pred_signs, n_pred = self.ada_aug_buf.tolist()
-
-            self.r_t_stat = pred_signs / n_pred
-
-            if self.r_t_stat > self.ada_aug_target:
-                sign = 1
-
-            else:
-                sign = -1
-
-            self.ada_aug_p += sign * n_pred / self.ada_aug_len
-            self.ada_aug_p = min(1, max(0, self.ada_aug_p))
-            self.ada_aug_buf.mul_(0)
-            self.ada_update = 0
-
-        return self.ada_aug_p
-
-
-SYM6 = (
-    0.015404109327027373,
-    0.0034907120842174702,
-    -0.11799011114819057,
-    -0.048311742585633,
-    0.4910559419267466,
-    0.787641141030194,
-    0.3379294217276218,
-    -0.07263752278646252,
-    -0.021060292512300564,
-    0.04472490177066578,
-    0.0017677118642428036,
-    -0.007800708325034148,
-)
-
-
-def translate_mat(t_x, t_y, device="cpu"):
-    batch = t_x.shape[0]
-
-    mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
-    translate = torch.stack((t_x, t_y), 1)
-    mat[:, :2, 2] = translate
-
-    return mat
-
-
-def rotate_mat(theta, device="cpu"):
-    batch = theta.shape[0]
-
-    mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
-    sin_t = torch.sin(theta)
-    cos_t = torch.cos(theta)
-    rot = torch.stack((cos_t, -sin_t, sin_t, cos_t), 1).view(batch, 2, 2)
-    mat[:, :2, :2] = rot
-
-    return mat
-
-
-def scale_mat(s_x, s_y, device="cpu"):
-    batch = s_x.shape[0]
-
-    mat = torch.eye(3, device=device).unsqueeze(0).repeat(batch, 1, 1)
-    mat[:, 0, 0] = s_x
-    mat[:, 1, 1] = s_y
-
-    return mat
-
-
-def translate3d_mat(t_x, t_y, t_z):
-    batch = t_x.shape[0]
-
-    mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
-    translate = torch.stack((t_x, t_y, t_z), 1)
-    mat[:, :3, 3] = translate
-
-    return mat
-
-
-def rotate3d_mat(axis, theta):
-    batch = theta.shape[0]
-
-    u_x, u_y, u_z = axis
-
-    eye = torch.eye(3).unsqueeze(0)
-    cross = torch.tensor([(0, -u_z, u_y), (u_z, 0, -u_x), (-u_y, u_x, 0)]).unsqueeze(0)
-    outer = torch.tensor(axis)
-    outer = (outer.unsqueeze(1) * outer).unsqueeze(0)
-
-    sin_t = torch.sin(theta).view(-1, 1, 1)
-    cos_t = torch.cos(theta).view(-1, 1, 1)
-
-    rot = cos_t * eye + sin_t * cross + (1 - cos_t) * outer
-
-    eye_4 = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
-    eye_4[:, :3, :3] = rot
-
-    return eye_4
-
-
-def scale3d_mat(s_x, s_y, s_z):
-    batch = s_x.shape[0]
-
-    mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
-    mat[:, 0, 0] = s_x
-    mat[:, 1, 1] = s_y
-    mat[:, 2, 2] = s_z
-
-    return mat
-
-
-def luma_flip_mat(axis, i):
-    batch = i.shape[0]
-
-    eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
-    axis = torch.tensor(axis + (0,))
-    flip = 2 * torch.ger(axis, axis) * i.view(-1, 1, 1)
-
-    return eye - flip
-
-
-def saturation_mat(axis, i):
-    batch = i.shape[0]
-
-    eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
-    axis = torch.tensor(axis + (0,))
-    axis = torch.ger(axis, axis)
-    saturate = axis + (eye - axis) * i.view(-1, 1, 1)
-
-    return saturate
-
-
-def lognormal_sample(size, mean=0, std=1, device="cpu"):
-    return torch.empty(size, device=device).log_normal_(mean=mean, std=std)
-
-
-def category_sample(size, categories, device="cpu"):
-    category = torch.tensor(categories, device=device)
-    sample = torch.randint(high=len(categories), size=(size,), device=device)
-
-    return category[sample]
-
-
-def uniform_sample(size, low, high, device="cpu"):
-    return torch.empty(size, device=device).uniform_(low, high)
-
-
-def normal_sample(size, mean=0, std=1, device="cpu"):
-    return torch.empty(size, device=device).normal_(mean, std)
-
-
-def bernoulli_sample(size, p, device="cpu"):
-    return torch.empty(size, device=device).bernoulli_(p)
-
-
-def random_mat_apply(p, transform, prev, eye, device="cpu"):
-    size = transform.shape[0]
-    select = bernoulli_sample(size, p, device=device).view(size, 1, 1)
-    select_transform = select * transform + (1 - select) * eye
-
-    return select_transform @ prev
-
-
-def sample_affine(p, size, height, width, device="cpu"):
-    G = torch.eye(3, device=device).unsqueeze(0).repeat(size, 1, 1)
-    eye = G
-
-    # flip
-    param = category_sample(size, (0, 1))
-    Gc = scale_mat(1 - 2.0 * param, torch.ones(size), device=device)
-    G = random_mat_apply(p, Gc, G, eye, device=device)
-    # print('flip', G, scale_mat(1 - 2.0 * param, torch.ones(size)), sep='\n')
-
-    # 90 rotate
-    param = category_sample(size, (0, 3))
-    Gc = rotate_mat(-math.pi / 2 * param, device=device)
-    G = random_mat_apply(p, Gc, G, eye, device=device)
-    # print('90 rotate', G, rotate_mat(-math.pi / 2 * param), sep='\n')
-
-    # integer translate
-    param = uniform_sample((2, size), -0.125, 0.125)
-    param_height = torch.round(param[0] * height)
-    param_width = torch.round(param[1] * width)
-    Gc = translate_mat(param_width, param_height, device=device)
-    G = random_mat_apply(p, Gc, G, eye, device=device)
-    # print('integer translate', G, translate_mat(param_width, param_height), sep='\n')
-
-    # isotropic scale
-    param = lognormal_sample(size, std=0.2 * math.log(2))
-    Gc = scale_mat(param, param, device=device)
-    G = random_mat_apply(p, Gc, G, eye, device=device)
-    # print('isotropic scale', G, scale_mat(param, param), sep='\n')
-
-    p_rot = 1 - math.sqrt(1 - p)
-
-    # pre-rotate
-    param = uniform_sample(size, -math.pi, math.pi)
-    Gc = rotate_mat(-param, device=device)
-    G = random_mat_apply(p_rot, Gc, G, eye, device=device)
-    # print('pre-rotate', G, rotate_mat(-param), sep='\n')
-
-    # anisotropic scale
-    param = lognormal_sample(size, std=0.2 * math.log(2))
-    Gc = scale_mat(param, 1 / param, device=device)
-    G = random_mat_apply(p, Gc, G, eye, device=device)
-    # print('anisotropic scale', G, scale_mat(param, 1 / param), sep='\n')
-
-    # post-rotate
-    param = uniform_sample(size, -math.pi, math.pi)
-    Gc = rotate_mat(-param, device=device)
-    G = random_mat_apply(p_rot, Gc, G, eye, device=device)
-    # print('post-rotate', G, rotate_mat(-param), sep='\n')
-
-    # fractional translate
-    param = normal_sample((2, size), std=0.125)
-    Gc = translate_mat(param[1] * width, param[0] * height, device=device)
-    G = random_mat_apply(p, Gc, G, eye, device=device)
-    # print('fractional translate', G, translate_mat(param, param), sep='\n')
-
-    return G
-
-
-def sample_color(p, size):
-    C = torch.eye(4).unsqueeze(0).repeat(size, 1, 1)
-    eye = C
-    axis_val = 1 / math.sqrt(3)
-    axis = (axis_val, axis_val, axis_val)
-
-    # brightness
-    param = normal_sample(size, std=0.2)
-    Cc = translate3d_mat(param, param, param)
-    C = random_mat_apply(p, Cc, C, eye)
-
-    # contrast
-    param = lognormal_sample(size, std=0.5 * math.log(2))
-    Cc = scale3d_mat(param, param, param)
-    C = random_mat_apply(p, Cc, C, eye)
-
-    # luma flip
-    param = category_sample(size, (0, 1))
-    Cc = luma_flip_mat(axis, param)
-    C = random_mat_apply(p, Cc, C, eye)
-
-    # hue rotation
-    param = uniform_sample(size, -math.pi, math.pi)
-    Cc = rotate3d_mat(axis, param)
-    C = random_mat_apply(p, Cc, C, eye)
-
-    # saturation
-    param = lognormal_sample(size, std=1 * math.log(2))
-    Cc = saturation_mat(axis, param)
-    C = random_mat_apply(p, Cc, C, eye)
-
-    return C
-
-
-def make_grid(shape, x0, x1, y0, y1, device):
-    n, c, h, w = shape
-    grid = torch.empty(n, h, w, 3, device=device)
-    grid[:, :, :, 0] = torch.linspace(x0, x1, w, device=device)
-    grid[:, :, :, 1] = torch.linspace(y0, y1, h, device=device).unsqueeze(-1)
-    grid[:, :, :, 2] = 1
-
-    return grid
-
-
-def affine_grid(grid, mat):
-    n, h, w, _ = grid.shape
-    return (grid.view(n, h * w, 3) @ mat.transpose(1, 2)).view(n, h, w, 2)
-
-
-def get_padding(G, height, width, kernel_size):
-    device = G.device
-
-    cx = (width - 1) / 2
-    cy = (height - 1) / 2
-    cp = torch.tensor(
-        [(-cx, -cy, 1), (cx, -cy, 1), (cx, cy, 1), (-cx, cy, 1)], device=device
-    )
-    cp = G @ cp.T
-
-    pad_k = kernel_size // 4
-
-    pad = cp[:, :2, :].permute(1, 0, 2).flatten(1)
-    pad = torch.cat((-pad, pad)).max(1).values
-    pad = pad + torch.tensor([pad_k * 2 - cx, pad_k * 2 - cy] * 2, device=device)
-    pad = pad.max(torch.tensor([0, 0] * 2, device=device))
-    pad = pad.min(torch.tensor([width - 1, height - 1] * 2, device=device))
-
-    pad_x1, pad_y1, pad_x2, pad_y2 = pad.ceil().to(torch.int32)
-
-    return pad_x1, pad_x2, pad_y1, pad_y2
-
-
-def try_sample_affine_and_pad(img, p, kernel_size, G=None):
-    batch, _, height, width = img.shape
-
-    G_try = G
-
-    if G is None:
-        G_try = torch.inverse(sample_affine(p, batch, height, width))
-
-    pad_x1, pad_x2, pad_y1, pad_y2 = get_padding(G_try, height, width, kernel_size)
-
-    img_pad = F.pad(img, (pad_x1, pad_x2, pad_y1, pad_y2), mode="reflect")
-
-    return img_pad, G_try, (pad_x1, pad_x2, pad_y1, pad_y2)
-
-
-class GridSampleForward(autograd.Function):
-    @staticmethod
-    def forward(ctx, input, grid):
-        out = F.grid_sample(
-            input, grid, mode="bilinear", padding_mode="zeros", align_corners=False
-        )
-        ctx.save_for_backward(input, grid)
-
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        input, grid = ctx.saved_tensors
-        grad_input, grad_grid = GridSampleBackward.apply(grad_output, input, grid)
-
-        return grad_input, grad_grid
-
-
-class GridSampleBackward(autograd.Function):
-    @staticmethod
-    def forward(ctx, grad_output, input, grid):
-        op = torch._C._jit_get_operation("aten::grid_sampler_2d_backward")
-        grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False)
-        ctx.save_for_backward(grid)
-
-        return grad_input, grad_grid
-
-    @staticmethod
-    def backward(ctx, grad_grad_input, grad_grad_grid):
-        (grid,) = ctx.saved_tensors
-        grad_grad_output = None
-
-        if ctx.needs_input_grad[0]:
-            grad_grad_output = GridSampleForward.apply(grad_grad_input, grid)
-
-        return grad_grad_output, None, None
-
-
-grid_sample = GridSampleForward.apply
-
-
-def scale_mat_single(s_x, s_y):
-    return torch.tensor(((s_x, 0, 0), (0, s_y, 0), (0, 0, 1)), dtype=torch.float32)
-
-
-def translate_mat_single(t_x, t_y):
-    return torch.tensor(((1, 0, t_x), (0, 1, t_y), (0, 0, 1)), dtype=torch.float32)
-
-
-def random_apply_affine(img, p, G=None, antialiasing_kernel=SYM6):
-    kernel = antialiasing_kernel
-    len_k = len(kernel)
-
-    kernel = torch.as_tensor(kernel).to(img)
-    # kernel = torch.ger(kernel, kernel).to(img)
-    kernel_flip = torch.flip(kernel, (0,))
-
-    img_pad, G, (pad_x1, pad_x2, pad_y1, pad_y2) = try_sample_affine_and_pad(
-        img, p, len_k, G
-    )
-
-    G_inv = (
-        translate_mat_single((pad_x1 - pad_x2).item() / 2, (pad_y1 - pad_y2).item() / 2)
-        @ G
-    )
-    up_pad = (
-        (len_k + 2 - 1) // 2,
-        (len_k - 2) // 2,
-        (len_k + 2 - 1) // 2,
-        (len_k - 2) // 2,
-    )
-    img_2x = upfirdn2d(img_pad, kernel.unsqueeze(0), up=(2, 1), pad=(*up_pad[:2], 0, 0))
-    img_2x = upfirdn2d(img_2x, kernel.unsqueeze(1), up=(1, 2), pad=(0, 0, *up_pad[2:]))
-    G_inv = scale_mat_single(2, 2) @ G_inv @ scale_mat_single(1 / 2, 1 / 2)
-    G_inv = translate_mat_single(-0.5, -0.5) @ G_inv @ translate_mat_single(0.5, 0.5)
-    batch_size, channel, height, width = img.shape
-    pad_k = len_k // 4
-    shape = (batch_size, channel, (height + pad_k * 2) * 2, (width + pad_k * 2) * 2)
-    G_inv = (
-        scale_mat_single(2 / img_2x.shape[3], 2 / img_2x.shape[2])
-        @ G_inv
-        @ scale_mat_single(1 / (2 / shape[3]), 1 / (2 / shape[2]))
-    )
-    grid = F.affine_grid(G_inv[:, :2, :].to(img_2x), shape, align_corners=False)
-    img_affine = grid_sample(img_2x, grid)
-    d_p = -pad_k * 2
-    down_pad = (
-        d_p + (len_k - 2 + 1) // 2,
-        d_p + (len_k - 2) // 2,
-        d_p + (len_k - 2 + 1) // 2,
-        d_p + (len_k - 2) // 2,
-    )
-    img_down = upfirdn2d(
-        img_affine, kernel_flip.unsqueeze(0), down=(2, 1), pad=(*down_pad[:2], 0, 0)
-    )
-    img_down = upfirdn2d(
-        img_down, kernel_flip.unsqueeze(1), down=(1, 2), pad=(0, 0, *down_pad[2:])
-    )
-
-    return img_down, G
-
-
-def apply_color(img, mat):
-    batch = img.shape[0]
-    img = img.permute(0, 2, 3, 1)
-    mat_mul = mat[:, :3, :3].transpose(1, 2).view(batch, 1, 3, 3)
-    mat_add = mat[:, :3, 3].view(batch, 1, 1, 3)
-    img = img @ mat_mul + mat_add
-    img = img.permute(0, 3, 1, 2)
-
-    return img
-
-
-def random_apply_color(img, p, C=None):
-    if C is None:
-        C = sample_color(p, img.shape[0])
-
-    img = apply_color(img, C.to(img))
-
-    return img, C
-
-
-def augment(img, p, transform_matrix=(None, None)):
-    img, G = random_apply_affine(img, p, transform_matrix[0])
-    img, C = random_apply_color(img, p, transform_matrix[1])
-
-    return img, (G, C)

op/__init__.py

@@ -1,2 +0,0 @@
-from .fused_act import FusedLeakyReLU, fused_leaky_relu
-from .upfirdn2d import upfirdn2d

op/conv2d_gradfix.py

@@ -1,227 +0,0 @@
-import contextlib
-import warnings
-
-import torch
-from torch import autograd
-from torch.nn import functional as F
-
-enabled = True
-weight_gradients_disabled = False
-
-
-@contextlib.contextmanager
-def no_weight_gradients():
-    global weight_gradients_disabled
-
-    old = weight_gradients_disabled
-    weight_gradients_disabled = True
-    yield
-    weight_gradients_disabled = old
-
-
-def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
-    if could_use_op(input):
-        return conv2d_gradfix(
-            transpose=False,
-            weight_shape=weight.shape,
-            stride=stride,
-            padding=padding,
-            output_padding=0,
-            dilation=dilation,
-            groups=groups,
-        ).apply(input, weight, bias)
-
-    return F.conv2d(
-        input=input,
-        weight=weight,
-        bias=bias,
-        stride=stride,
-        padding=padding,
-        dilation=dilation,
-        groups=groups,
-    )
-
-
-def conv_transpose2d(
-    input,
-    weight,
-    bias=None,
-    stride=1,
-    padding=0,
-    output_padding=0,
-    groups=1,
-    dilation=1,
-):
-    if could_use_op(input):
-        return conv2d_gradfix(
-            transpose=True,
-            weight_shape=weight.shape,
-            stride=stride,
-            padding=padding,
-            output_padding=output_padding,
-            groups=groups,
-            dilation=dilation,
-        ).apply(input, weight, bias)
-
-    return F.conv_transpose2d(
-        input=input,
-        weight=weight,
-        bias=bias,
-        stride=stride,
-        padding=padding,
-        output_padding=output_padding,
-        dilation=dilation,
-        groups=groups,
-    )
-
-
-def could_use_op(input):
-    if (not enabled) or (not torch.backends.cudnn.enabled):
-        return False
-
-    if input.device.type != "cuda":
-        return False
-
-    if any(torch.__version__.startswith(x) for x in ["1.7.", "1.8."]):
-        return True
-
-    warnings.warn(
-        f"conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d()."
-    )
-
-    return False
-
-
-def ensure_tuple(xs, ndim):
-    xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
-
-    return xs
-
-
-conv2d_gradfix_cache = dict()
-
-
-def conv2d_gradfix(
-    transpose, weight_shape, stride, padding, output_padding, dilation, groups
-):
-    ndim = 2
-    weight_shape = tuple(weight_shape)
-    stride = ensure_tuple(stride, ndim)
-    padding = ensure_tuple(padding, ndim)
-    output_padding = ensure_tuple(output_padding, ndim)
-    dilation = ensure_tuple(dilation, ndim)
-
-    key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
-    if key in conv2d_gradfix_cache:
-        return conv2d_gradfix_cache[key]
-
-    common_kwargs = dict(
-        stride=stride, padding=padding, dilation=dilation, groups=groups
-    )
-
-    def calc_output_padding(input_shape, output_shape):
-        if transpose:
-            return [0, 0]
-
-        return [
-            input_shape[i + 2]
-            - (output_shape[i + 2] - 1) * stride[i]
-            - (1 - 2 * padding[i])
-            - dilation[i] * (weight_shape[i + 2] - 1)
-            for i in range(ndim)
-        ]
-
-    class Conv2d(autograd.Function):
-        @staticmethod
-        def forward(ctx, input, weight, bias):
-            if not transpose:
-                out = F.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)
-
-            else:
-                out = F.conv_transpose2d(
-                    input=input,
-                    weight=weight,
-                    bias=bias,
-                    output_padding=output_padding,
-                    **common_kwargs,
-                )
-
-            ctx.save_for_backward(input, weight)
-
-            return out
-
-        @staticmethod
-        def backward(ctx, grad_output):
-            input, weight = ctx.saved_tensors
-            grad_input, grad_weight, grad_bias = None, None, None
-
-            if ctx.needs_input_grad[0]:
-                p = calc_output_padding(
-                    input_shape=input.shape, output_shape=grad_output.shape
-                )
-                grad_input = conv2d_gradfix(
-                    transpose=(not transpose),
-                    weight_shape=weight_shape,
-                    output_padding=p,
-                    **common_kwargs,
-                ).apply(grad_output, weight, None)
-
-            if ctx.needs_input_grad[1] and not weight_gradients_disabled:
-                grad_weight = Conv2dGradWeight.apply(grad_output, input)
-
-            if ctx.needs_input_grad[2]:
-                grad_bias = grad_output.sum((0, 2, 3))
-
-            return grad_input, grad_weight, grad_bias
-
-    class Conv2dGradWeight(autograd.Function):
-        @staticmethod
-        def forward(ctx, grad_output, input):
-            op = torch._C._jit_get_operation(
-                "aten::cudnn_convolution_backward_weight"
-                if not transpose
-                else "aten::cudnn_convolution_transpose_backward_weight"
-            )
-            flags = [
-                torch.backends.cudnn.benchmark,
-                torch.backends.cudnn.deterministic,
-                torch.backends.cudnn.allow_tf32,
-            ]
-            grad_weight = op(
-                weight_shape,
-                grad_output,
-                input,
-                padding,
-                stride,
-                dilation,
-                groups,
-                *flags,
-            )
-            ctx.save_for_backward(grad_output, input)
-
-            return grad_weight
-
-        @staticmethod
-        def backward(ctx, grad_grad_weight):
-            grad_output, input = ctx.saved_tensors
-            grad_grad_output, grad_grad_input = None, None
-
-            if ctx.needs_input_grad[0]:
-                grad_grad_output = Conv2d.apply(input, grad_grad_weight, None)
-
-            if ctx.needs_input_grad[1]:
-                p = calc_output_padding(
-                    input_shape=input.shape, output_shape=grad_output.shape
-                )
-                grad_grad_input = conv2d_gradfix(
-                    transpose=(not transpose),
-                    weight_shape=weight_shape,
-                    output_padding=p,
-                    **common_kwargs,
-                ).apply(grad_output, grad_grad_weight, None)
-
-            return grad_grad_output, grad_grad_input
-
-    conv2d_gradfix_cache[key] = Conv2d
-
-    return Conv2d

op/fused_act.py

@@ -1,127 +0,0 @@
-import os
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.autograd import Function
-from torch.utils.cpp_extension import load
-
-
-module_path = os.path.dirname(__file__)
-fused = load(
-    "fused",
-    sources=[
-        os.path.join(module_path, "fused_bias_act.cpp"),
-        os.path.join(module_path, "fused_bias_act_kernel.cu"),
-    ],
-)
-
-
-class FusedLeakyReLUFunctionBackward(Function):
-    @staticmethod
-    def forward(ctx, grad_output, out, bias, negative_slope, scale):
-        ctx.save_for_backward(out)
-        ctx.negative_slope = negative_slope
-        ctx.scale = scale
-
-        empty = grad_output.new_empty(0)
-
-        grad_input = fused.fused_bias_act(
-            grad_output.contiguous(), empty, out, 3, 1, negative_slope, scale
-        )
-
-        dim = [0]
-
-        if grad_input.ndim > 2:
-            dim += list(range(2, grad_input.ndim))
-
-        if bias:
-            grad_bias = grad_input.sum(dim).detach()
-
-        else:
-            grad_bias = empty
-
-        return grad_input, grad_bias
-
-    @staticmethod
-    def backward(ctx, gradgrad_input, gradgrad_bias):
-        out, = ctx.saved_tensors
-        gradgrad_out = fused.fused_bias_act(
-            gradgrad_input.contiguous(),
-            gradgrad_bias,
-            out,
-            3,
-            1,
-            ctx.negative_slope,
-            ctx.scale,
-        )
-
-        return gradgrad_out, None, None, None, None
-
-
-class FusedLeakyReLUFunction(Function):
-    @staticmethod
-    def forward(ctx, input, bias, negative_slope, scale):
-        empty = input.new_empty(0)
-
-        ctx.bias = bias is not None
-
-        if bias is None:
-            bias = empty
-
-        out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
-        ctx.save_for_backward(out)
-        ctx.negative_slope = negative_slope
-        ctx.scale = scale
-
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        out, = ctx.saved_tensors
-
-        grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(
-            grad_output, out, ctx.bias, ctx.negative_slope, ctx.scale
-        )
-
-        if not ctx.bias:
-            grad_bias = None
-
-        return grad_input, grad_bias, None, None
-
-
-class FusedLeakyReLU(nn.Module):
-    def __init__(self, channel, bias=True, negative_slope=0.2, scale=2 ** 0.5):
-        super().__init__()
-
-        if bias:
-            self.bias = nn.Parameter(torch.zeros(channel))
-
-        else:
-            self.bias = None
-
-        self.negative_slope = negative_slope
-        self.scale = scale
-
-    def forward(self, input):
-        return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
-
-
-def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2 ** 0.5):
-    if input.device.type == "cpu":
-        if bias is not None:
-            rest_dim = [1] * (input.ndim - bias.ndim - 1)
-            return (
-                F.leaky_relu(
-                    input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2
-                )
-                * scale
-            )
-
-        else:
-            return F.leaky_relu(input, negative_slope=0.2) * scale
-
-    else:
-        return FusedLeakyReLUFunction.apply(
-            input.contiguous(), bias, negative_slope, scale
-        )

op/fused_bias_act.cpp

@@ -1,32 +0,0 @@
-
-#include <ATen/ATen.h>
-#include <torch/extension.h>
-
-torch::Tensor fused_bias_act_op(const torch::Tensor &input,
-                                const torch::Tensor &bias,
-                                const torch::Tensor &refer, int act, int grad,
-                                float alpha, float scale);
-
-#define CHECK_CUDA(x)                                                          \
-  TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x)                                                    \
-  TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
-#define CHECK_INPUT(x)                                                         \
-  CHECK_CUDA(x);                                                               \
-  CHECK_CONTIGUOUS(x)
-
-torch::Tensor fused_bias_act(const torch::Tensor &input,
-                             const torch::Tensor &bias,
-                             const torch::Tensor &refer, int act, int grad,
-                             float alpha, float scale) {
-  CHECK_INPUT(input);
-  CHECK_INPUT(bias);
-
-  at::DeviceGuard guard(input.device());
-
-  return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale);
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)");
-}

op/fused_bias_act_kernel.cu

@@ -1,105 +0,0 @@
-// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
-//
-// This work is made available under the Nvidia Source Code License-NC.
-// To view a copy of this license, visit
-// https://nvlabs.github.io/stylegan2/license.html
-
-#include <torch/types.h>
-
-#include <ATen/ATen.h>
-#include <ATen/AccumulateType.h>
-#include <ATen/cuda/CUDAApplyUtils.cuh>
-#include <ATen/cuda/CUDAContext.h>
-
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-template <typename scalar_t>
-static __global__ void
-fused_bias_act_kernel(scalar_t *out, const scalar_t *p_x, const scalar_t *p_b,
-                      const scalar_t *p_ref, int act, int grad, scalar_t alpha,
-                      scalar_t scale, int loop_x, int size_x, int step_b,
-                      int size_b, int use_bias, int use_ref) {
-  int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;
-
-  scalar_t zero = 0.0;
-
-  for (int loop_idx = 0; loop_idx < loop_x && xi < size_x;
-       loop_idx++, xi += blockDim.x) {
-    scalar_t x = p_x[xi];
-
-    if (use_bias) {
-      x += p_b[(xi / step_b) % size_b];
-    }
-
-    scalar_t ref = use_ref ? p_ref[xi] : zero;
-
-    scalar_t y;
-
-    switch (act * 10 + grad) {
-    default:
-    case 10:
-      y = x;
-      break;
-    case 11:
-      y = x;
-      break;
-    case 12:
-      y = 0.0;
-      break;
-
-    case 30:
-      y = (x > 0.0) ? x : x * alpha;
-      break;
-    case 31:
-      y = (ref > 0.0) ? x : x * alpha;
-      break;
-    case 32:
-      y = 0.0;
-      break;
-    }
-
-    out[xi] = y * scale;
-  }
-}
-
-torch::Tensor fused_bias_act_op(const torch::Tensor &input,
-                                const torch::Tensor &bias,
-                                const torch::Tensor &refer, int act, int grad,
-                                float alpha, float scale) {
-  int curDevice = -1;
-  cudaGetDevice(&curDevice);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  auto x = input.contiguous();
-  auto b = bias.contiguous();
-  auto ref = refer.contiguous();
-
-  int use_bias = b.numel() ? 1 : 0;
-  int use_ref = ref.numel() ? 1 : 0;
-
-  int size_x = x.numel();
-  int size_b = b.numel();
-  int step_b = 1;
-
-  for (int i = 1 + 1; i < x.dim(); i++) {
-    step_b *= x.size(i);
-  }
-
-  int loop_x = 4;
-  int block_size = 4 * 32;
-  int grid_size = (size_x - 1) / (loop_x * block_size) + 1;
-
-  auto y = torch::empty_like(x);
-
-  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-      x.scalar_type(), "fused_bias_act_kernel", [&] {
-        fused_bias_act_kernel<scalar_t><<<grid_size, block_size, 0, stream>>>(
-            y.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(),
-            b.data_ptr<scalar_t>(), ref.data_ptr<scalar_t>(), act, grad, alpha,
-            scale, loop_x, size_x, step_b, size_b, use_bias, use_ref);
-      });
-
-  return y;
-}

op/upfirdn2d.cpp

@@ -1,31 +0,0 @@
-#include <ATen/ATen.h>
-#include <torch/extension.h>
-
-torch::Tensor upfirdn2d_op(const torch::Tensor &input,
-                           const torch::Tensor &kernel, int up_x, int up_y,
-                           int down_x, int down_y, int pad_x0, int pad_x1,
-                           int pad_y0, int pad_y1);
-
-#define CHECK_CUDA(x)                                                          \
-  TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x)                                                    \
-  TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
-#define CHECK_INPUT(x)                                                         \
-  CHECK_CUDA(x);                                                               \
-  CHECK_CONTIGUOUS(x)
-
-torch::Tensor upfirdn2d(const torch::Tensor &input, const torch::Tensor &kernel,
-                        int up_x, int up_y, int down_x, int down_y, int pad_x0,
-                        int pad_x1, int pad_y0, int pad_y1) {
-  CHECK_INPUT(input);
-  CHECK_INPUT(kernel);
-
-  at::DeviceGuard guard(input.device());
-
-  return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1,
-                      pad_y0, pad_y1);
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)");
-}

op/upfirdn2d.py

@@ -1,209 +0,0 @@
-from collections import abc
-import os
-
-import torch
-from torch.nn import functional as F
-from torch.autograd import Function
-from torch.utils.cpp_extension import load
-
-
-module_path = os.path.dirname(__file__)
-upfirdn2d_op = load(
-    "upfirdn2d",
-    sources=[
-        os.path.join(module_path, "upfirdn2d.cpp"),
-        os.path.join(module_path, "upfirdn2d_kernel.cu"),
-    ],
-)
-
-
-class UpFirDn2dBackward(Function):
-    @staticmethod
-    def forward(
-        ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
-    ):
-
-        up_x, up_y = up
-        down_x, down_y = down
-        g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad
-
-        grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)
-
-        grad_input = upfirdn2d_op.upfirdn2d(
-            grad_output,
-            grad_kernel,
-            down_x,
-            down_y,
-            up_x,
-            up_y,
-            g_pad_x0,
-            g_pad_x1,
-            g_pad_y0,
-            g_pad_y1,
-        )
-        grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])
-
-        ctx.save_for_backward(kernel)
-
-        pad_x0, pad_x1, pad_y0, pad_y1 = pad
-
-        ctx.up_x = up_x
-        ctx.up_y = up_y
-        ctx.down_x = down_x
-        ctx.down_y = down_y
-        ctx.pad_x0 = pad_x0
-        ctx.pad_x1 = pad_x1
-        ctx.pad_y0 = pad_y0
-        ctx.pad_y1 = pad_y1
-        ctx.in_size = in_size
-        ctx.out_size = out_size
-
-        return grad_input
-
-    @staticmethod
-    def backward(ctx, gradgrad_input):
-        kernel, = ctx.saved_tensors
-
-        gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)
-
-        gradgrad_out = upfirdn2d_op.upfirdn2d(
-            gradgrad_input,
-            kernel,
-            ctx.up_x,
-            ctx.up_y,
-            ctx.down_x,
-            ctx.down_y,
-            ctx.pad_x0,
-            ctx.pad_x1,
-            ctx.pad_y0,
-            ctx.pad_y1,
-        )
-        # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])
-        gradgrad_out = gradgrad_out.view(
-            ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
-        )
-
-        return gradgrad_out, None, None, None, None, None, None, None, None
-
-
-class UpFirDn2d(Function):
-    @staticmethod
-    def forward(ctx, input, kernel, up, down, pad):
-        up_x, up_y = up
-        down_x, down_y = down
-        pad_x0, pad_x1, pad_y0, pad_y1 = pad
-
-        kernel_h, kernel_w = kernel.shape
-        batch, channel, in_h, in_w = input.shape
-        ctx.in_size = input.shape
-
-        input = input.reshape(-1, in_h, in_w, 1)
-
-        ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))
-
-        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y
-        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x
-        ctx.out_size = (out_h, out_w)
-
-        ctx.up = (up_x, up_y)
-        ctx.down = (down_x, down_y)
-        ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)
-
-        g_pad_x0 = kernel_w - pad_x0 - 1
-        g_pad_y0 = kernel_h - pad_y0 - 1
-        g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
-        g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1
-
-        ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)
-
-        out = upfirdn2d_op.upfirdn2d(
-            input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
-        )
-        # out = out.view(major, out_h, out_w, minor)
-        out = out.view(-1, channel, out_h, out_w)
-
-        return out
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        kernel, grad_kernel = ctx.saved_tensors
-
-        grad_input = None
-
-        if ctx.needs_input_grad[0]:
-            grad_input = UpFirDn2dBackward.apply(
-                grad_output,
-                kernel,
-                grad_kernel,
-                ctx.up,
-                ctx.down,
-                ctx.pad,
-                ctx.g_pad,
-                ctx.in_size,
-                ctx.out_size,
-            )
-
-        return grad_input, None, None, None, None
-
-
-def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
-    if not isinstance(up, abc.Iterable):
-        up = (up, up)
-
-    if not isinstance(down, abc.Iterable):
-        down = (down, down)
-
-    if len(pad) == 2:
-        pad = (pad[0], pad[1], pad[0], pad[1])
-
-    if input.device.type == "cpu":
-        out = upfirdn2d_native(input, kernel, *up, *down, *pad)
-
-    else:
-        out = UpFirDn2d.apply(input, kernel, up, down, pad)
-
-    return out
-
-
-def upfirdn2d_native(
-    input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
-):
-    _, channel, in_h, in_w = input.shape
-    input = input.reshape(-1, in_h, in_w, 1)
-
-    _, in_h, in_w, minor = input.shape
-    kernel_h, kernel_w = kernel.shape
-
-    out = input.view(-1, in_h, 1, in_w, 1, minor)
-    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
-    out = out.view(-1, in_h * up_y, in_w * up_x, minor)
-
-    out = F.pad(
-        out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
-    )
-    out = out[
-        :,
-        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
-        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
-        :,
-    ]
-
-    out = out.permute(0, 3, 1, 2)
-    out = out.reshape(
-        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
-    )
-    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
-    out = F.conv2d(out, w)
-    out = out.reshape(
-        -1,
-        minor,
-        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
-        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
-    )
-    out = out.permute(0, 2, 3, 1)
-    out = out[:, ::down_y, ::down_x, :]
-
-    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y
-    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x
-
-    return out.view(-1, channel, out_h, out_w)

op/upfirdn2d_kernel.cu

@@ -1,369 +0,0 @@
-// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
-//
-// This work is made available under the Nvidia Source Code License-NC.
-// To view a copy of this license, visit
-// https://nvlabs.github.io/stylegan2/license.html
-
-#include <torch/types.h>
-
-#include <ATen/ATen.h>
-#include <ATen/AccumulateType.h>
-#include <ATen/cuda/CUDAApplyUtils.cuh>
-#include <ATen/cuda/CUDAContext.h>
-
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-static __host__ __device__ __forceinline__ int floor_div(int a, int b) {
-  int c = a / b;
-
-  if (c * b > a) {
-    c--;
-  }
-
-  return c;
-}
-
-struct UpFirDn2DKernelParams {
-  int up_x;
-  int up_y;
-  int down_x;
-  int down_y;
-  int pad_x0;
-  int pad_x1;
-  int pad_y0;
-  int pad_y1;
-
-  int major_dim;
-  int in_h;
-  int in_w;
-  int minor_dim;
-  int kernel_h;
-  int kernel_w;
-  int out_h;
-  int out_w;
-  int loop_major;
-  int loop_x;
-};
-
-template <typename scalar_t>
-__global__ void upfirdn2d_kernel_large(scalar_t *out, const scalar_t *input,
-                                       const scalar_t *kernel,
-                                       const UpFirDn2DKernelParams p) {
-  int minor_idx = blockIdx.x * blockDim.x + threadIdx.x;
-  int out_y = minor_idx / p.minor_dim;
-  minor_idx -= out_y * p.minor_dim;
-  int out_x_base = blockIdx.y * p.loop_x * blockDim.y + threadIdx.y;
-  int major_idx_base = blockIdx.z * p.loop_major;
-
-  if (out_x_base >= p.out_w || out_y >= p.out_h ||
-      major_idx_base >= p.major_dim) {
-    return;
-  }
-
-  int mid_y = out_y * p.down_y + p.up_y - 1 - p.pad_y0;
-  int in_y = min(max(floor_div(mid_y, p.up_y), 0), p.in_h);
-  int h = min(max(floor_div(mid_y + p.kernel_h, p.up_y), 0), p.in_h) - in_y;
-  int kernel_y = mid_y + p.kernel_h - (in_y + 1) * p.up_y;
-
-  for (int loop_major = 0, major_idx = major_idx_base;
-       loop_major < p.loop_major && major_idx < p.major_dim;
-       loop_major++, major_idx++) {
-    for (int loop_x = 0, out_x = out_x_base;
-         loop_x < p.loop_x && out_x < p.out_w; loop_x++, out_x += blockDim.y) {
-      int mid_x = out_x * p.down_x + p.up_x - 1 - p.pad_x0;
-      int in_x = min(max(floor_div(mid_x, p.up_x), 0), p.in_w);
-      int w = min(max(floor_div(mid_x + p.kernel_w, p.up_x), 0), p.in_w) - in_x;
-      int kernel_x = mid_x + p.kernel_w - (in_x + 1) * p.up_x;
-
-      const scalar_t *x_p =
-          &input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim +
-                 minor_idx];
-      const scalar_t *k_p = &kernel[kernel_y * p.kernel_w + kernel_x];
-      int x_px = p.minor_dim;
-      int k_px = -p.up_x;
-      int x_py = p.in_w * p.minor_dim;
-      int k_py = -p.up_y * p.kernel_w;
-
-      scalar_t v = 0.0f;
-
-      for (int y = 0; y < h; y++) {
-        for (int x = 0; x < w; x++) {
-          v += static_cast<scalar_t>(*x_p) * static_cast<scalar_t>(*k_p);
-          x_p += x_px;
-          k_p += k_px;
-        }
-
-        x_p += x_py - w * x_px;
-        k_p += k_py - w * k_px;
-      }
-
-      out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +
-          minor_idx] = v;
-    }
-  }
-}
-
-template <typename scalar_t, int up_x, int up_y, int down_x, int down_y,
-          int kernel_h, int kernel_w, int tile_out_h, int tile_out_w>
-__global__ void upfirdn2d_kernel(scalar_t *out, const scalar_t *input,
-                                 const scalar_t *kernel,
-                                 const UpFirDn2DKernelParams p) {
-  const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1;
-  const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1;
-
-  __shared__ volatile float sk[kernel_h][kernel_w];
-  __shared__ volatile float sx[tile_in_h][tile_in_w];
-
-  int minor_idx = blockIdx.x;
-  int tile_out_y = minor_idx / p.minor_dim;
-  minor_idx -= tile_out_y * p.minor_dim;
-  tile_out_y *= tile_out_h;
-  int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w;
-  int major_idx_base = blockIdx.z * p.loop_major;
-
-  if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h |
-      major_idx_base >= p.major_dim) {
-    return;
-  }
-
-  for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w;
-       tap_idx += blockDim.x) {
-    int ky = tap_idx / kernel_w;
-    int kx = tap_idx - ky * kernel_w;
-    scalar_t v = 0.0;
-
-    if (kx < p.kernel_w & ky < p.kernel_h) {
-      v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)];
-    }
-
-    sk[ky][kx] = v;
-  }
-
-  for (int loop_major = 0, major_idx = major_idx_base;
-       loop_major < p.loop_major & major_idx < p.major_dim;
-       loop_major++, major_idx++) {
-    for (int loop_x = 0, tile_out_x = tile_out_x_base;
-         loop_x < p.loop_x & tile_out_x < p.out_w;
-         loop_x++, tile_out_x += tile_out_w) {
-      int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0;
-      int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0;
-      int tile_in_x = floor_div(tile_mid_x, up_x);
-      int tile_in_y = floor_div(tile_mid_y, up_y);
-
-      __syncthreads();
-
-      for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w;
-           in_idx += blockDim.x) {
-        int rel_in_y = in_idx / tile_in_w;
-        int rel_in_x = in_idx - rel_in_y * tile_in_w;
-        int in_x = rel_in_x + tile_in_x;
-        int in_y = rel_in_y + tile_in_y;
-
-        scalar_t v = 0.0;
-
-        if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) {
-          v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) *
-                        p.minor_dim +
-                    minor_idx];
-        }
-
-        sx[rel_in_y][rel_in_x] = v;
-      }
-
-      __syncthreads();
-      for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w;
-           out_idx += blockDim.x) {
-        int rel_out_y = out_idx / tile_out_w;
-        int rel_out_x = out_idx - rel_out_y * tile_out_w;
-        int out_x = rel_out_x + tile_out_x;
-        int out_y = rel_out_y + tile_out_y;
-
-        int mid_x = tile_mid_x + rel_out_x * down_x;
-        int mid_y = tile_mid_y + rel_out_y * down_y;
-        int in_x = floor_div(mid_x, up_x);
-        int in_y = floor_div(mid_y, up_y);
-        int rel_in_x = in_x - tile_in_x;
-        int rel_in_y = in_y - tile_in_y;
-        int kernel_x = (in_x + 1) * up_x - mid_x - 1;
-        int kernel_y = (in_y + 1) * up_y - mid_y - 1;
-
-        scalar_t v = 0.0;
-
-#pragma unroll
-        for (int y = 0; y < kernel_h / up_y; y++)
-#pragma unroll
-          for (int x = 0; x < kernel_w / up_x; x++)
-            v += sx[rel_in_y + y][rel_in_x + x] *
-                 sk[kernel_y + y * up_y][kernel_x + x * up_x];
-
-        if (out_x < p.out_w & out_y < p.out_h) {
-          out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +
-              minor_idx] = v;
-        }
-      }
-    }
-  }
-}
-
-torch::Tensor upfirdn2d_op(const torch::Tensor &input,
-                           const torch::Tensor &kernel, int up_x, int up_y,
-                           int down_x, int down_y, int pad_x0, int pad_x1,
-                           int pad_y0, int pad_y1) {
-  int curDevice = -1;
-  cudaGetDevice(&curDevice);
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  UpFirDn2DKernelParams p;
-
-  auto x = input.contiguous();
-  auto k = kernel.contiguous();
-
-  p.major_dim = x.size(0);
-  p.in_h = x.size(1);
-  p.in_w = x.size(2);
-  p.minor_dim = x.size(3);
-  p.kernel_h = k.size(0);
-  p.kernel_w = k.size(1);
-  p.up_x = up_x;
-  p.up_y = up_y;
-  p.down_x = down_x;
-  p.down_y = down_y;
-  p.pad_x0 = pad_x0;
-  p.pad_x1 = pad_x1;
-  p.pad_y0 = pad_y0;
-  p.pad_y1 = pad_y1;
-
-  p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) /
-            p.down_y;
-  p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) /
-            p.down_x;
-
-  auto out =
-      at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options());
-
-  int mode = -1;
-
-  int tile_out_h = -1;
-  int tile_out_w = -1;
-
-  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&
-      p.kernel_h <= 4 && p.kernel_w <= 4) {
-    mode = 1;
-    tile_out_h = 16;
-    tile_out_w = 64;
-  }
-
-  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&
-      p.kernel_h <= 3 && p.kernel_w <= 3) {
-    mode = 2;
-    tile_out_h = 16;
-    tile_out_w = 64;
-  }
-
-  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&
-      p.kernel_h <= 4 && p.kernel_w <= 4) {
-    mode = 3;
-    tile_out_h = 16;
-    tile_out_w = 64;
-  }
-
-  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&
-      p.kernel_h <= 2 && p.kernel_w <= 2) {
-    mode = 4;
-    tile_out_h = 16;
-    tile_out_w = 64;
-  }
-
-  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&
-      p.kernel_h <= 4 && p.kernel_w <= 4) {
-    mode = 5;
-    tile_out_h = 8;
-    tile_out_w = 32;
-  }
-
-  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&
-      p.kernel_h <= 2 && p.kernel_w <= 2) {
-    mode = 6;
-    tile_out_h = 8;
-    tile_out_w = 32;
-  }
-
-  dim3 block_size;
-  dim3 grid_size;
-
-  if (tile_out_h > 0 && tile_out_w > 0) {
-    p.loop_major = (p.major_dim - 1) / 16384 + 1;
-    p.loop_x = 1;
-    block_size = dim3(32 * 8, 1, 1);
-    grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim,
-                     (p.out_w - 1) / (p.loop_x * tile_out_w) + 1,
-                     (p.major_dim - 1) / p.loop_major + 1);
-  } else {
-    p.loop_major = (p.major_dim - 1) / 16384 + 1;
-    p.loop_x = 4;
-    block_size = dim3(4, 32, 1);
-    grid_size = dim3((p.out_h * p.minor_dim - 1) / block_size.x + 1,
-                     (p.out_w - 1) / (p.loop_x * block_size.y) + 1,
-                     (p.major_dim - 1) / p.loop_major + 1);
-  }
-
-  AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] {
-    switch (mode) {
-    case 1:
-      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 4, 4, 16, 64>
-          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
-                                                 x.data_ptr<scalar_t>(),
-                                                 k.data_ptr<scalar_t>(), p);
-
-      break;
-
-    case 2:
-      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 3, 3, 16, 64>
-          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
-                                                 x.data_ptr<scalar_t>(),
-                                                 k.data_ptr<scalar_t>(), p);
-
-      break;
-
-    case 3:
-      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 4, 4, 16, 64>
-          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
-                                                 x.data_ptr<scalar_t>(),
-                                                 k.data_ptr<scalar_t>(), p);
-
-      break;
-
-    case 4:
-      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 2, 2, 16, 64>
-          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
-                                                 x.data_ptr<scalar_t>(),
-                                                 k.data_ptr<scalar_t>(), p);
-
-      break;
-
-    case 5:
-      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>
-          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
-                                                 x.data_ptr<scalar_t>(),
-                                                 k.data_ptr<scalar_t>(), p);
-
-      break;
-
-    case 6:
-      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>
-          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
-                                                 x.data_ptr<scalar_t>(),
-                                                 k.data_ptr<scalar_t>(), p);
-
-      break;
-
-    default:
-      upfirdn2d_kernel_large<scalar_t><<<grid_size, block_size, 0, stream>>>(
-          out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(),
-          k.data_ptr<scalar_t>(), p);
-    }
-  });
-
-  return out;
-}

ppl.py

@@ -1,130 +0,0 @@
-import argparse
-
-import torch
-from torch.nn import functional as F
-import numpy as np
-from tqdm import tqdm
-
-import lpips
-from model import Generator
-
-
-def normalize(x):
-    return x / torch.sqrt(x.pow(2).sum(-1, keepdim=True))
-
-
-def slerp(a, b, t):
-    a = normalize(a)
-    b = normalize(b)
-    d = (a * b).sum(-1, keepdim=True)
-    p = t * torch.acos(d)
-    c = normalize(b - d * a)
-    d = a * torch.cos(p) + c * torch.sin(p)
-
-    return normalize(d)
-
-
-def lerp(a, b, t):
-    return a + (b - a) * t
-
-
-if __name__ == "__main__":
-    device = "cuda"
-
-    parser = argparse.ArgumentParser(description="Perceptual Path Length calculator")
-
-    parser.add_argument(
-        "--space", choices=["z", "w"], help="space that PPL calculated with"
-    )
-    parser.add_argument(
-        "--batch", type=int, default=64, help="batch size for the models"
-    )
-    parser.add_argument(
-        "--n_sample",
-        type=int,
-        default=5000,
-        help="number of the samples for calculating PPL",
-    )
-    parser.add_argument(
-        "--size", type=int, default=256, help="output image sizes of the generator"
-    )
-    parser.add_argument(
-        "--eps", type=float, default=1e-4, help="epsilon for numerical stability"
-    )
-    parser.add_argument(
-        "--crop", action="store_true", help="apply center crop to the images"
-    )
-    parser.add_argument(
-        "--sampling",
-        default="end",
-        choices=["end", "full"],
-        help="set endpoint sampling method",
-    )
-    parser.add_argument(
-        "ckpt", metavar="CHECKPOINT", help="path to the model checkpoints"
-    )
-
-    args = parser.parse_args()
-
-    latent_dim = 512
-
-    ckpt = torch.load(args.ckpt)
-
-    g = Generator(args.size, latent_dim, 8).to(device)
-    g.load_state_dict(ckpt["g_ema"])
-    g.eval()
-
-    percept = lpips.PerceptualLoss(
-        model="net-lin", net="vgg", use_gpu=device.startswith("cuda")
-    )
-
-    distances = []
-
-    n_batch = args.n_sample // args.batch
-    resid = args.n_sample - (n_batch * args.batch)
-    batch_sizes = [args.batch] * n_batch + [resid]
-
-    with torch.no_grad():
-        for batch in tqdm(batch_sizes):
-            noise = g.make_noise()
-
-            inputs = torch.randn([batch * 2, latent_dim], device=device)
-            if args.sampling == "full":
-                lerp_t = torch.rand(batch, device=device)
-            else:
-                lerp_t = torch.zeros(batch, device=device)
-
-            if args.space == "w":
-                latent = g.get_latent(inputs)
-                latent_t0, latent_t1 = latent[::2], latent[1::2]
-                latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None])
-                latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None] + args.eps)
-                latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)
-
-            image, _ = g([latent_e], input_is_latent=True, noise=noise)
-
-            if args.crop:
-                c = image.shape[2] // 8
-                image = image[:, :, c * 3 : c * 7, c * 2 : c * 6]
-
-            factor = image.shape[2] // 256
-
-            if factor > 1:
-                image = F.interpolate(
-                    image, size=(256, 256), mode="bilinear", align_corners=False
-                )
-
-            dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (
-                args.eps ** 2
-            )
-            distances.append(dist.to("cpu").numpy())
-
-    distances = np.concatenate(distances, 0)
-
-    lo = np.percentile(distances, 1, interpolation="lower")
-    hi = np.percentile(distances, 99, interpolation="higher")
-    filtered_dist = np.extract(
-        np.logical_and(lo <= distances, distances <= hi), distances
-    )
-
-    print("ppl:", filtered_dist.mean())

prepare_data.py

@@ -1,101 +0,0 @@
-import argparse
-from io import BytesIO
-import multiprocessing
-from functools import partial
-
-from PIL import Image
-import lmdb
-from tqdm import tqdm
-from torchvision import datasets
-from torchvision.transforms import functional as trans_fn
-
-
-def resize_and_convert(img, size, resample, quality=100):
-    img = trans_fn.resize(img, size, resample)
-    img = trans_fn.center_crop(img, size)
-    buffer = BytesIO()
-    img.save(buffer, format="jpeg", quality=quality)
-    val = buffer.getvalue()
-
-    return val
-
-
-def resize_multiple(
-    img, sizes=(128, 256, 512, 1024), resample=Image.LANCZOS, quality=100
-):
-    imgs = []
-
-    for size in sizes:
-        imgs.append(resize_and_convert(img, size, resample, quality))
-
-    return imgs
-
-
-def resize_worker(img_file, sizes, resample):
-    i, file = img_file
-    img = Image.open(file)
-    img = img.convert("RGB")
-    out = resize_multiple(img, sizes=sizes, resample=resample)
-
-    return i, out
-
-
-def prepare(
-    env, dataset, n_worker, sizes=(128, 256, 512, 1024), resample=Image.LANCZOS
-):
-    resize_fn = partial(resize_worker, sizes=sizes, resample=resample)
-
-    files = sorted(dataset.imgs, key=lambda x: x[0])
-    files = [(i, file) for i, (file, label) in enumerate(files)]
-    total = 0
-
-    with multiprocessing.Pool(n_worker) as pool:
-        for i, imgs in tqdm(pool.imap_unordered(resize_fn, files)):
-            for size, img in zip(sizes, imgs):
-                key = f"{size}-{str(i).zfill(5)}".encode("utf-8")
-
-                with env.begin(write=True) as txn:
-                    txn.put(key, img)
-
-            total += 1
-
-        with env.begin(write=True) as txn:
-            txn.put("length".encode("utf-8"), str(total).encode("utf-8"))
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Preprocess images for model training")
-    parser.add_argument("--out", type=str, help="filename of the result lmdb dataset")
-    parser.add_argument(
-        "--size",
-        type=str,
-        default="128,256,512,1024",
-        help="resolutions of images for the dataset",
-    )
-    parser.add_argument(
-        "--n_worker",
-        type=int,
-        default=8,
-        help="number of workers for preparing dataset",
-    )
-    parser.add_argument(
-        "--resample",
-        type=str,
-        default="lanczos",
-        help="resampling methods for resizing images",
-    )
-    parser.add_argument("path", type=str, help="path to the image dataset")
-
-    args = parser.parse_args()
-
-    resample_map = {"lanczos": Image.LANCZOS, "bilinear": Image.BILINEAR}
-    resample = resample_map[args.resample]
-
-    sizes = [int(s.strip()) for s in args.size.split(",")]
-
-    print(f"Make dataset of image sizes:", ", ".join(str(s) for s in sizes))
-
-    imgset = datasets.ImageFolder(args.path)
-
-    with lmdb.open(args.out, map_size=1024 ** 4, readahead=False) as env:
-        prepare(env, imgset, args.n_worker, sizes=sizes, resample=resample)

projector.py

@@ -1,248 +0,0 @@
-import argparse
-import math
-import os
-
-import torch
-from torch import optim
-from torch.nn import functional as F
-from torchvision import transforms
-from PIL import Image
-from tqdm import tqdm
-
-import lpips
-from model import Generator
-
-
-def noise_regularize(noises):
-    loss = 0
-
-    for noise in noises:
-        size = noise.shape[2]
-
-        while True:
-            loss = (
-                loss
-                + (noise * torch.roll(noise, shifts=1, dims=3)).mean().pow(2)
-                + (noise * torch.roll(noise, shifts=1, dims=2)).mean().pow(2)
-            )
-
-            if size <= 8:
-                break
-
-            noise = noise.reshape([-1, 1, size // 2, 2, size // 2, 2])
-            noise = noise.mean([3, 5])
-            size //= 2
-
-    return loss
-
-
-def noise_normalize_(noises):
-    for noise in noises:
-        mean = noise.mean()
-        std = noise.std()
-
-        noise.data.add_(-mean).div_(std)
-
-
-def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05):
-    lr_ramp = min(1, (1 - t) / rampdown)
-    lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi)
-    lr_ramp = lr_ramp * min(1, t / rampup)
-
-    return initial_lr * lr_ramp
-
-
-def latent_noise(latent, strength):
-    noise = torch.randn_like(latent) * strength
-
-    return latent + noise
-
-
-def make_image(tensor):
-    return (
-        tensor.detach()
-        .clamp_(min=-1, max=1)
-        .add(1)
-        .div_(2)
-        .mul(255)
-        .type(torch.uint8)
-        .permute(0, 2, 3, 1)
-        .to("cpu")
-        .numpy()
-    )
-
-
-if __name__ == "__main__":
-    device = "cuda"
-
-    parser = argparse.ArgumentParser(
-        description="Image projector to the generator latent spaces"
-    )
-    parser.add_argument(
-        "--ckpt", type=str, required=True, help="path to the model checkpoint"
-    )
-    parser.add_argument(
-        "--size", type=int, default=256, help="output image sizes of the generator"
-    )
-    parser.add_argument(
-        "--lr_rampup",
-        type=float,
-        default=0.05,
-        help="duration of the learning rate warmup",
-    )
-    parser.add_argument(
-        "--lr_rampdown",
-        type=float,
-        default=0.25,
-        help="duration of the learning rate decay",
-    )
-    parser.add_argument("--lr", type=float, default=0.1, help="learning rate")
-    parser.add_argument(
-        "--noise", type=float, default=0.05, help="strength of the noise level"
-    )
-    parser.add_argument(
-        "--noise_ramp",
-        type=float,
-        default=0.75,
-        help="duration of the noise level decay",
-    )
-    parser.add_argument("--step", type=int, default=1000, help="optimize iterations")
-    parser.add_argument(
-        "--noise_regularize",
-        type=float,
-        default=1e5,
-        help="weight of the noise regularization",
-    )
-    parser.add_argument("--mse", type=float, default=0, help="weight of the mse loss")
-    parser.add_argument(
-        "--w_plus",
-        action="store_true",
-        help="allow to use distinct latent codes to each layers",
-    )
-    parser.add_argument(
-        "files", metavar="FILES", nargs="+", help="path to image files to be projected"
-    )
-
-    args = parser.parse_args()
-
-    n_mean_latent = 10000
-
-    resize = min(args.size, 256)
-
-    transform = transforms.Compose(
-        [
-            transforms.Resize(resize),
-            transforms.CenterCrop(resize),
-            transforms.ToTensor(),
-            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
-        ]
-    )
-
-    imgs = []
-
-    for imgfile in args.files:
-        img = transform(Image.open(imgfile).convert("RGB"))
-        imgs.append(img)
-
-    imgs = torch.stack(imgs, 0).to(device)
-
-    g_ema = Generator(args.size, 512, 8)
-    g_ema.load_state_dict(torch.load(args.ckpt)["g_ema"], strict=False)
-    g_ema.eval()
-    g_ema = g_ema.to(device)
-
-    with torch.no_grad():
-        noise_sample = torch.randn(n_mean_latent, 512, device=device)
-        latent_out = g_ema.style(noise_sample)
-
-        latent_mean = latent_out.mean(0)
-        latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5
-
-    percept = lpips.PerceptualLoss(
-        model="net-lin", net="vgg", use_gpu=device.startswith("cuda")
-    )
-
-    noises_single = g_ema.make_noise()
-    noises = []
-    for noise in noises_single:
-        noises.append(noise.repeat(imgs.shape[0], 1, 1, 1).normal_())
-
-    latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(imgs.shape[0], 1)
-
-    if args.w_plus:
-        latent_in = latent_in.unsqueeze(1).repeat(1, g_ema.n_latent, 1)
-
-    latent_in.requires_grad = True
-
-    for noise in noises:
-        noise.requires_grad = True
-
-    optimizer = optim.Adam([latent_in] + noises, lr=args.lr)
-
-    pbar = tqdm(range(args.step))
-    latent_path = []
-
-    for i in pbar:
-        t = i / args.step
-        lr = get_lr(t, args.lr)
-        optimizer.param_groups[0]["lr"] = lr
-        noise_strength = latent_std * args.noise * max(0, 1 - t / args.noise_ramp) ** 2
-        latent_n = latent_noise(latent_in, noise_strength.item())
-
-        img_gen, _ = g_ema([latent_n], input_is_latent=True, noise=noises)
-
-        batch, channel, height, width = img_gen.shape
-
-        if height > 256:
-            factor = height // 256
-
-            img_gen = img_gen.reshape(
-                batch, channel, height // factor, factor, width // factor, factor
-            )
-            img_gen = img_gen.mean([3, 5])
-
-        p_loss = percept(img_gen, imgs).sum()
-        n_loss = noise_regularize(noises)
-        mse_loss = F.mse_loss(img_gen, imgs)
-
-        loss = p_loss + args.noise_regularize * n_loss + args.mse * mse_loss
-
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        noise_normalize_(noises)
-
-        if (i + 1) % 100 == 0:
-            latent_path.append(latent_in.detach().clone())
-
-        pbar.set_description(
-            (
-                f"perceptual: {p_loss.item():.4f}; noise regularize: {n_loss.item():.4f};"
-                f" mse: {mse_loss.item():.4f}; lr: {lr:.4f}"
-            )
-        )
-
-    img_gen, _ = g_ema([latent_path[-1]], input_is_latent=True, noise=noises)
-
-    filename = os.path.splitext(os.path.basename(args.files[0]))[0] + ".pt"
-
-    img_ar = make_image(img_gen)
-
-    result_file = {}
-    for i, input_name in enumerate(args.files):
-        noise_single = []
-        for noise in noises:
-            noise_single.append(noise[i : i + 1])
-
-        result_file[input_name] = {
-            "img": img_gen[i],
-            "latent": latent_in[i],
-            "noise": noise_single,
-        }
-
-        img_name = os.path.splitext(os.path.basename(input_name))[0] + "-project.png"
-        pil_img = Image.fromarray(img_ar[i])
-        pil_img.save(img_name)
-
-    torch.save(result_file, filename)

sample/.gitignore

@@ -1 +0,0 @@
-*.png

swagan.py

@@ -1,440 +0,0 @@
-import math
-import random
-import functools
-import operator
-
-import torch
-from torch import nn
-from torch.nn import functional as F
-from torch.autograd import Function
-
-from op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d, conv2d_gradfix
-from model import (
-    ModulatedConv2d,
-    StyledConv,
-    ConstantInput,
-    PixelNorm,
-    Upsample,
-    Downsample,
-    Blur,
-    EqualLinear,
-    ConvLayer,
-)
-
-
-def get_haar_wavelet(in_channels):
-    haar_wav_l = 1 / (2 ** 0.5) * torch.ones(1, 2)
-    haar_wav_h = 1 / (2 ** 0.5) * torch.ones(1, 2)
-    haar_wav_h[0, 0] = -1 * haar_wav_h[0, 0]
-
-    haar_wav_ll = haar_wav_l.T * haar_wav_l
-    haar_wav_lh = haar_wav_h.T * haar_wav_l
-    haar_wav_hl = haar_wav_l.T * haar_wav_h
-    haar_wav_hh = haar_wav_h.T * haar_wav_h
-
-    return haar_wav_ll, haar_wav_lh, haar_wav_hl, haar_wav_hh
-
-
-def dwt_init(x):
-    x01 = x[:, :, 0::2, :] / 2
-    x02 = x[:, :, 1::2, :] / 2
-    x1 = x01[:, :, :, 0::2]
-    x2 = x02[:, :, :, 0::2]
-    x3 = x01[:, :, :, 1::2]
-    x4 = x02[:, :, :, 1::2]
-    x_LL = x1 + x2 + x3 + x4
-    x_HL = -x1 - x2 + x3 + x4
-    x_LH = -x1 + x2 - x3 + x4
-    x_HH = x1 - x2 - x3 + x4
-
-    return torch.cat((x_LL, x_HL, x_LH, x_HH), 1)
-
-
-def iwt_init(x):
-    r = 2
-    in_batch, in_channel, in_height, in_width = x.size()
-    # print([in_batch, in_channel, in_height, in_width])
-    out_batch, out_channel, out_height, out_width = (
-        in_batch,
-        int(in_channel / (r ** 2)),
-        r * in_height,
-        r * in_width,
-    )
-    x1 = x[:, 0:out_channel, :, :] / 2
-    x2 = x[:, out_channel : out_channel * 2, :, :] / 2
-    x3 = x[:, out_channel * 2 : out_channel * 3, :, :] / 2
-    x4 = x[:, out_channel * 3 : out_channel * 4, :, :] / 2
-
-    h = torch.zeros([out_batch, out_channel, out_height, out_width]).float().cuda()
-
-    h[:, :, 0::2, 0::2] = x1 - x2 - x3 + x4
-    h[:, :, 1::2, 0::2] = x1 - x2 + x3 - x4
-    h[:, :, 0::2, 1::2] = x1 + x2 - x3 - x4
-    h[:, :, 1::2, 1::2] = x1 + x2 + x3 + x4
-
-    return h
-
-
-class HaarTransform(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-
-        ll, lh, hl, hh = get_haar_wavelet(in_channels)
-
-        self.register_buffer("ll", ll)
-        self.register_buffer("lh", lh)
-        self.register_buffer("hl", hl)
-        self.register_buffer("hh", hh)
-
-    def forward(self, input):
-        ll = upfirdn2d(input, self.ll, down=2)
-        lh = upfirdn2d(input, self.lh, down=2)
-        hl = upfirdn2d(input, self.hl, down=2)
-        hh = upfirdn2d(input, self.hh, down=2)
-
-        return torch.cat((ll, lh, hl, hh), 1)
-
-
-class InverseHaarTransform(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-
-        ll, lh, hl, hh = get_haar_wavelet(in_channels)
-
-        self.register_buffer("ll", ll)
-        self.register_buffer("lh", -lh)
-        self.register_buffer("hl", -hl)
-        self.register_buffer("hh", hh)
-
-    def forward(self, input):
-        ll, lh, hl, hh = input.chunk(4, 1)
-        ll = upfirdn2d(ll, self.ll, up=2, pad=(1, 0, 1, 0))
-        lh = upfirdn2d(lh, self.lh, up=2, pad=(1, 0, 1, 0))
-        hl = upfirdn2d(hl, self.hl, up=2, pad=(1, 0, 1, 0))
-        hh = upfirdn2d(hh, self.hh, up=2, pad=(1, 0, 1, 0))
-
-        return ll + lh + hl + hh
-
-
-class ToRGB(nn.Module):
-    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        if upsample:
-            self.iwt = InverseHaarTransform(3)
-            self.upsample = Upsample(blur_kernel)
-            self.dwt = HaarTransform(3)
-
-        self.conv = ModulatedConv2d(in_channel, 3 * 4, 1, style_dim, demodulate=False)
-        self.bias = nn.Parameter(torch.zeros(1, 3 * 4, 1, 1))
-
-    def forward(self, input, style, skip=None):
-        out = self.conv(input, style)
-        out = out + self.bias
-
-        if skip is not None:
-            skip = self.iwt(skip)
-            skip = self.upsample(skip)
-            skip = self.dwt(skip)
-
-            out = out + skip
-
-        return out
-
-
-class Generator(nn.Module):
-    def __init__(
-        self,
-        size,
-        style_dim,
-        n_mlp,
-        channel_multiplier=2,
-        blur_kernel=[1, 3, 3, 1],
-        lr_mlp=0.01,
-    ):
-        super().__init__()
-
-        self.size = size
-
-        self.style_dim = style_dim
-
-        layers = [PixelNorm()]
-
-        for i in range(n_mlp):
-            layers.append(
-                EqualLinear(
-                    style_dim, style_dim, lr_mul=lr_mlp, activation="fused_lrelu"
-                )
-            )
-
-        self.style = nn.Sequential(*layers)
-
-        self.channels = {
-            4: 512,
-            8: 512,
-            16: 512,
-            32: 512,
-            64: 256 * channel_multiplier,
-            128: 128 * channel_multiplier,
-            256: 64 * channel_multiplier,
-            512: 32 * channel_multiplier,
-            1024: 16 * channel_multiplier,
-        }
-
-        self.input = ConstantInput(self.channels[4])
-        self.conv1 = StyledConv(
-            self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel
-        )
-        self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)
-
-        self.log_size = int(math.log(size, 2)) - 1
-        self.num_layers = (self.log_size - 2) * 2 + 1
-
-        self.convs = nn.ModuleList()
-        self.upsamples = nn.ModuleList()
-        self.to_rgbs = nn.ModuleList()
-        self.noises = nn.Module()
-
-        in_channel = self.channels[4]
-
-        for layer_idx in range(self.num_layers):
-            res = (layer_idx + 5) // 2
-            shape = [1, 1, 2 ** res, 2 ** res]
-            self.noises.register_buffer(f"noise_{layer_idx}", torch.randn(*shape))
-
-        for i in range(3, self.log_size + 1):
-            out_channel = self.channels[2 ** i]
-
-            self.convs.append(
-                StyledConv(
-                    in_channel,
-                    out_channel,
-                    3,
-                    style_dim,
-                    upsample=True,
-                    blur_kernel=blur_kernel,
-                )
-            )
-
-            self.convs.append(
-                StyledConv(
-                    out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel
-                )
-            )
-
-            self.to_rgbs.append(ToRGB(out_channel, style_dim))
-
-            in_channel = out_channel
-
-        self.iwt = InverseHaarTransform(3)
-
-        self.n_latent = self.log_size * 2 - 2
-
-    def make_noise(self):
-        device = self.input.input.device
-
-        noises = [torch.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]
-
-        for i in range(3, self.log_size + 1):
-            for _ in range(2):
-                noises.append(torch.randn(1, 1, 2 ** i, 2 ** i, device=device))
-
-        return noises
-
-    def mean_latent(self, n_latent):
-        latent_in = torch.randn(
-            n_latent, self.style_dim, device=self.input.input.device
-        )
-        latent = self.style(latent_in).mean(0, keepdim=True)
-
-        return latent
-
-    def get_latent(self, input):
-        return self.style(input)
-
-    def forward(
-        self,
-        styles,
-        return_latents=False,
-        inject_index=None,
-        truncation=1,
-        truncation_latent=None,
-        input_is_latent=False,
-        noise=None,
-        randomize_noise=True,
-    ):
-        if not input_is_latent:
-            styles = [self.style(s) for s in styles]
-
-        if noise is None:
-            if randomize_noise:
-                noise = [None] * self.num_layers
-            else:
-                noise = [
-                    getattr(self.noises, f"noise_{i}") for i in range(self.num_layers)
-                ]
-
-        if truncation < 1:
-            style_t = []
-
-            for style in styles:
-                style_t.append(
-                    truncation_latent + truncation * (style - truncation_latent)
-                )
-
-            styles = style_t
-
-        if len(styles) < 2:
-            inject_index = self.n_latent
-
-            if styles[0].ndim < 3:
-                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
-
-            else:
-                latent = styles[0]
-
-        else:
-            if inject_index is None:
-                inject_index = random.randint(1, self.n_latent - 1)
-
-            latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
-            latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)
-
-            latent = torch.cat([latent, latent2], 1)
-
-        out = self.input(latent)
-        out = self.conv1(out, latent[:, 0], noise=noise[0])
-
-        skip = self.to_rgb1(out, latent[:, 1])
-
-        i = 1
-        for conv1, conv2, noise1, noise2, to_rgb in zip(
-            self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
-        ):
-            out = conv1(out, latent[:, i], noise=noise1)
-            out = conv2(out, latent[:, i + 1], noise=noise2)
-            skip = to_rgb(out, latent[:, i + 2], skip)
-
-            i += 2
-
-        image = self.iwt(skip)
-
-        if return_latents:
-            return image, latent
-
-        else:
-            return image, None
-
-
-class ConvBlock(nn.Module):
-    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        self.conv1 = ConvLayer(in_channel, in_channel, 3)
-        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
-
-    def forward(self, input):
-        out = self.conv1(input)
-        out = self.conv2(out)
-
-        return out
-
-
-class FromRGB(nn.Module):
-    def __init__(self, out_channel, downsample=True, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        self.downsample = downsample
-
-        if downsample:
-            self.iwt = InverseHaarTransform(3)
-            self.downsample = Downsample(blur_kernel)
-            self.dwt = HaarTransform(3)
-
-        self.conv = ConvLayer(3 * 4, out_channel, 3)
-
-    def forward(self, input, skip=None):
-        if self.downsample:
-            input = self.iwt(input)
-            input = self.downsample(input)
-            input = self.dwt(input)
-
-        out = self.conv(input)
-
-        if skip is not None:
-            out = out + skip
-
-        return input, out
-
-
-class Discriminator(nn.Module):
-    def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
-        super().__init__()
-
-        channels = {
-            4: 512,
-            8: 512,
-            16: 512,
-            32: 512,
-            64: 256 * channel_multiplier,
-            128: 128 * channel_multiplier,
-            256: 64 * channel_multiplier,
-            512: 32 * channel_multiplier,
-            1024: 16 * channel_multiplier,
-        }
-
-        self.dwt = HaarTransform(3)
-
-        self.from_rgbs = nn.ModuleList()
-        self.convs = nn.ModuleList()
-
-        log_size = int(math.log(size, 2)) - 1
-
-        in_channel = channels[size]
-
-        for i in range(log_size, 2, -1):
-            out_channel = channels[2 ** (i - 1)]
-
-            self.from_rgbs.append(FromRGB(in_channel, downsample=i != log_size))
-            self.convs.append(ConvBlock(in_channel, out_channel, blur_kernel))
-
-            in_channel = out_channel
-
-        self.from_rgbs.append(FromRGB(channels[4]))
-
-        self.stddev_group = 4
-        self.stddev_feat = 1
-
-        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
-        self.final_linear = nn.Sequential(
-            EqualLinear(channels[4] * 4 * 4, channels[4], activation="fused_lrelu"),
-            EqualLinear(channels[4], 1),
-        )
-
-    def forward(self, input):
-        input = self.dwt(input)
-        out = None
-
-        for from_rgb, conv in zip(self.from_rgbs, self.convs):
-            input, out = from_rgb(input, out)
-            out = conv(out)
-
-        _, out = self.from_rgbs[-1](input, out)
-
-        batch, channel, height, width = out.shape
-        group = min(batch, self.stddev_group)
-        stddev = out.view(
-            group, -1, self.stddev_feat, channel // self.stddev_feat, height, width
-        )
-        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-8)
-        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
-        stddev = stddev.repeat(group, 1, height, width)
-        out = torch.cat([out, stddev], 1)
-
-        out = self.final_conv(out)
-
-        out = out.view(batch, -1)
-        out = self.final_linear(out)
-
-        return out
-

train.py

@@ -1,531 +0,0 @@
-import argparse
-import math
-import random
-import os
-
-import numpy as np
-import torch
-from torch import nn, autograd, optim
-from torch.nn import functional as F
-from torch.utils import data
-import torch.distributed as dist
-from torchvision import transforms, utils
-from tqdm import tqdm
-
-try:
-    import wandb
-
-except ImportError:
-    wandb = None
-
-
-from dataset import MultiResolutionDataset
-from distributed import (
-    get_rank,
-    synchronize,
-    reduce_loss_dict,
-    reduce_sum,
-    get_world_size,
-)
-from op import conv2d_gradfix
-from non_leaking import augment, AdaptiveAugment
-
-
-def data_sampler(dataset, shuffle, distributed):
-    if distributed:
-        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)
-
-    if shuffle:
-        return data.RandomSampler(dataset)
-
-    else:
-        return data.SequentialSampler(dataset)
-
-
-def requires_grad(model, flag=True):
-    for p in model.parameters():
-        p.requires_grad = flag
-
-
-def accumulate(model1, model2, decay=0.999):
-    par1 = dict(model1.named_parameters())
-    par2 = dict(model2.named_parameters())
-
-    for k in par1.keys():
-        par1[k].data.mul_(decay).add_(par2[k].data, alpha=1 - decay)
-
-
-def sample_data(loader):
-    while True:
-        for batch in loader:
-            yield batch
-
-
-def d_logistic_loss(real_pred, fake_pred):
-    real_loss = F.softplus(-real_pred)
-    fake_loss = F.softplus(fake_pred)
-
-    return real_loss.mean() + fake_loss.mean()
-
-
-def d_r1_loss(real_pred, real_img):
-    with conv2d_gradfix.no_weight_gradients():
-        grad_real, = autograd.grad(
-            outputs=real_pred.sum(), inputs=real_img, create_graph=True
-        )
-    grad_penalty = grad_real.pow(2).reshape(grad_real.shape[0], -1).sum(1).mean()
-
-    return grad_penalty
-
-
-def g_nonsaturating_loss(fake_pred):
-    loss = F.softplus(-fake_pred).mean()
-
-    return loss
-
-
-def g_path_regularize(fake_img, latents, mean_path_length, decay=0.01):
-    noise = torch.randn_like(fake_img) / math.sqrt(
-        fake_img.shape[2] * fake_img.shape[3]
-    )
-    grad, = autograd.grad(
-        outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True
-    )
-    path_lengths = torch.sqrt(grad.pow(2).sum(2).mean(1))
-
-    path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)
-
-    path_penalty = (path_lengths - path_mean).pow(2).mean()
-
-    return path_penalty, path_mean.detach(), path_lengths
-
-
-def make_noise(batch, latent_dim, n_noise, device):
-    if n_noise == 1:
-        return torch.randn(batch, latent_dim, device=device)
-
-    noises = torch.randn(n_noise, batch, latent_dim, device=device).unbind(0)
-
-    return noises
-
-
-def mixing_noise(batch, latent_dim, prob, device):
-    if prob > 0 and random.random() < prob:
-        return make_noise(batch, latent_dim, 2, device)
-
-    else:
-        return [make_noise(batch, latent_dim, 1, device)]
-
-
-def set_grad_none(model, targets):
-    for n, p in model.named_parameters():
-        if n in targets:
-            p.grad = None
-
-
-def train(args, loader, generator, discriminator, g_optim, d_optim, g_ema, device):
-    loader = sample_data(loader)
-
-    pbar = range(args.iter)
-
-    if get_rank() == 0:
-        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0.01)
-
-    mean_path_length = 0
-
-    d_loss_val = 0
-    r1_loss = torch.tensor(0.0, device=device)
-    g_loss_val = 0
-    path_loss = torch.tensor(0.0, device=device)
-    path_lengths = torch.tensor(0.0, device=device)
-    mean_path_length_avg = 0
-    loss_dict = {}
-
-    if args.distributed:
-        g_module = generator.module
-        d_module = discriminator.module
-
-    else:
-        g_module = generator
-        d_module = discriminator
-
-    accum = 0.5 ** (32 / (10 * 1000))
-    ada_aug_p = args.augment_p if args.augment_p > 0 else 0.0
-    r_t_stat = 0
-
-    if args.augment and args.augment_p == 0:
-        ada_augment = AdaptiveAugment(args.ada_target, args.ada_length, 8, device)
-
-    sample_z = torch.randn(args.n_sample, args.latent, device=device)
-
-    for idx in pbar:
-        i = idx + args.start_iter
-
-        if i > args.iter:
-            print("Done!")
-
-            break
-
-        real_img = next(loader)
-        real_img = real_img.to(device)
-
-        requires_grad(generator, False)
-        requires_grad(discriminator, True)
-
-        noise = mixing_noise(args.batch, args.latent, args.mixing, device)
-        fake_img, _ = generator(noise)
-
-        if args.augment:
-            real_img_aug, _ = augment(real_img, ada_aug_p)
-            fake_img, _ = augment(fake_img, ada_aug_p)
-
-        else:
-            real_img_aug = real_img
-
-        fake_pred = discriminator(fake_img)
-        real_pred = discriminator(real_img_aug)
-        d_loss = d_logistic_loss(real_pred, fake_pred)
-
-        loss_dict["d"] = d_loss
-        loss_dict["real_score"] = real_pred.mean()
-        loss_dict["fake_score"] = fake_pred.mean()
-
-        discriminator.zero_grad()
-        d_loss.backward()
-        d_optim.step()
-
-        if args.augment and args.augment_p == 0:
-            ada_aug_p = ada_augment.tune(real_pred)
-            r_t_stat = ada_augment.r_t_stat
-
-        d_regularize = i % args.d_reg_every == 0
-
-        if d_regularize:
-            real_img.requires_grad = True
-
-            if args.augment:
-                real_img_aug, _ = augment(real_img, ada_aug_p)
-
-            else:
-                real_img_aug = real_img
-
-            real_pred = discriminator(real_img_aug)
-            r1_loss = d_r1_loss(real_pred, real_img)
-
-            discriminator.zero_grad()
-            (args.r1 / 2 * r1_loss * args.d_reg_every + 0 * real_pred[0]).backward()
-
-            d_optim.step()
-
-        loss_dict["r1"] = r1_loss
-
-        requires_grad(generator, True)
-        requires_grad(discriminator, False)
-
-        noise = mixing_noise(args.batch, args.latent, args.mixing, device)
-        fake_img, _ = generator(noise)
-
-        if args.augment:
-            fake_img, _ = augment(fake_img, ada_aug_p)
-
-        fake_pred = discriminator(fake_img)
-        g_loss = g_nonsaturating_loss(fake_pred)
-
-        loss_dict["g"] = g_loss
-
-        generator.zero_grad()
-        g_loss.backward()
-        g_optim.step()
-
-        g_regularize = i % args.g_reg_every == 0
-
-        if g_regularize:
-            path_batch_size = max(1, args.batch // args.path_batch_shrink)
-            noise = mixing_noise(path_batch_size, args.latent, args.mixing, device)
-            fake_img, latents = generator(noise, return_latents=True)
-
-            path_loss, mean_path_length, path_lengths = g_path_regularize(
-                fake_img, latents, mean_path_length
-            )
-
-            generator.zero_grad()
-            weighted_path_loss = args.path_regularize * args.g_reg_every * path_loss
-
-            if args.path_batch_shrink:
-                weighted_path_loss += 0 * fake_img[0, 0, 0, 0]
-
-            weighted_path_loss.backward()
-
-            g_optim.step()
-
-            mean_path_length_avg = (
-                reduce_sum(mean_path_length).item() / get_world_size()
-            )
-
-        loss_dict["path"] = path_loss
-        loss_dict["path_length"] = path_lengths.mean()
-
-        accumulate(g_ema, g_module, accum)
-
-        loss_reduced = reduce_loss_dict(loss_dict)
-
-        d_loss_val = loss_reduced["d"].mean().item()
-        g_loss_val = loss_reduced["g"].mean().item()
-        r1_val = loss_reduced["r1"].mean().item()
-        path_loss_val = loss_reduced["path"].mean().item()
-        real_score_val = loss_reduced["real_score"].mean().item()
-        fake_score_val = loss_reduced["fake_score"].mean().item()
-        path_length_val = loss_reduced["path_length"].mean().item()
-
-        if get_rank() == 0:
-            pbar.set_description(
-                (
-                    f"d: {d_loss_val:.4f}; g: {g_loss_val:.4f}; r1: {r1_val:.4f}; "
-                    f"path: {path_loss_val:.4f}; mean path: {mean_path_length_avg:.4f}; "
-                    f"augment: {ada_aug_p:.4f}"
-                )
-            )
-
-            if wandb and args.wandb:
-                wandb.log(
-                    {
-                        "Generator": g_loss_val,
-                        "Discriminator": d_loss_val,
-                        "Augment": ada_aug_p,
-                        "Rt": r_t_stat,
-                        "R1": r1_val,
-                        "Path Length Regularization": path_loss_val,
-                        "Mean Path Length": mean_path_length,
-                        "Real Score": real_score_val,
-                        "Fake Score": fake_score_val,
-                        "Path Length": path_length_val,
-                    }
-                )
-
-            if i % 100 == 0:
-                with torch.no_grad():
-                    g_ema.eval()
-                    sample, _ = g_ema([sample_z])
-                    utils.save_image(
-                        sample,
-                        f"sample/{str(i).zfill(6)}.png",
-                        nrow=int(args.n_sample ** 0.5),
-                        normalize=True,
-                        range=(-1, 1),
-                    )
-
-            if i % 10000 == 0:
-                torch.save(
-                    {
-                        "g": g_module.state_dict(),
-                        "d": d_module.state_dict(),
-                        "g_ema": g_ema.state_dict(),
-                        "g_optim": g_optim.state_dict(),
-                        "d_optim": d_optim.state_dict(),
-                        "args": args,
-                        "ada_aug_p": ada_aug_p,
-                    },
-                    f"checkpoint/{str(i).zfill(6)}.pt",
-                )
-
-
-if __name__ == "__main__":
-    device = "cuda"
-
-    parser = argparse.ArgumentParser(description="StyleGAN2 trainer")
-
-    parser.add_argument("path", type=str, help="path to the lmdb dataset")
-    parser.add_argument('--arch', type=str, default='stylegan2', help='model architectures (stylegan2 | swagan)')
-    parser.add_argument(
-        "--iter", type=int, default=800000, help="total training iterations"
-    )
-    parser.add_argument(
-        "--batch", type=int, default=16, help="batch sizes for each gpus"
-    )
-    parser.add_argument(
-        "--n_sample",
-        type=int,
-        default=64,
-        help="number of the samples generated during training",
-    )
-    parser.add_argument(
-        "--size", type=int, default=256, help="image sizes for the model"
-    )
-    parser.add_argument(
-        "--r1", type=float, default=10, help="weight of the r1 regularization"
-    )
-    parser.add_argument(
-        "--path_regularize",
-        type=float,
-        default=2,
-        help="weight of the path length regularization",
-    )
-    parser.add_argument(
-        "--path_batch_shrink",
-        type=int,
-        default=2,
-        help="batch size reducing factor for the path length regularization (reduce memory consumption)",
-    )
-    parser.add_argument(
-        "--d_reg_every",
-        type=int,
-        default=16,
-        help="interval of the applying r1 regularization",
-    )
-    parser.add_argument(
-        "--g_reg_every",
-        type=int,
-        default=4,
-        help="interval of the applying path length regularization",
-    )
-    parser.add_argument(
-        "--mixing", type=float, default=0.9, help="probability of latent code mixing"
-    )
-    parser.add_argument(
-        "--ckpt",
-        type=str,
-        default=None,
-        help="path to the checkpoints to resume training",
-    )
-    parser.add_argument("--lr", type=float, default=0.002, help="learning rate")
-    parser.add_argument(
-        "--channel_multiplier",
-        type=int,
-        default=2,
-        help="channel multiplier factor for the model. config-f = 2, else = 1",
-    )
-    parser.add_argument(
-        "--wandb", action="store_true", help="use weights and biases logging"
-    )
-    parser.add_argument(
-        "--local_rank", type=int, default=0, help="local rank for distributed training"
-    )
-    parser.add_argument(
-        "--augment", action="store_true", help="apply non leaking augmentation"
-    )
-    parser.add_argument(
-        "--augment_p",
-        type=float,
-        default=0,
-        help="probability of applying augmentation. 0 = use adaptive augmentation",
-    )
-    parser.add_argument(
-        "--ada_target",
-        type=float,
-        default=0.6,
-        help="target augmentation probability for adaptive augmentation",
-    )
-    parser.add_argument(
-        "--ada_length",
-        type=int,
-        default=500 * 1000,
-        help="target duraing to reach augmentation probability for adaptive augmentation",
-    )
-    parser.add_argument(
-        "--ada_every",
-        type=int,
-        default=256,
-        help="probability update interval of the adaptive augmentation",
-    )
-
-    args = parser.parse_args()
-
-    n_gpu = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
-    args.distributed = n_gpu > 1
-
-    if args.distributed:
-        torch.cuda.set_device(args.local_rank)
-        torch.distributed.init_process_group(backend="nccl", init_method="env://")
-        synchronize()
-
-    args.latent = 512
-    args.n_mlp = 8
-
-    args.start_iter = 0
-
-    if args.arch == 'stylegan2':
-        from model import Generator, Discriminator
-
-    elif args.arch == 'swagan':
-        from swagan import Generator, Discriminator
-
-    generator = Generator(
-        args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier
-    ).to(device)
-    discriminator = Discriminator(
-        args.size, channel_multiplier=args.channel_multiplier
-    ).to(device)
-    g_ema = Generator(
-        args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier
-    ).to(device)
-    g_ema.eval()
-    accumulate(g_ema, generator, 0)
-
-    g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
-    d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)
-
-    g_optim = optim.Adam(
-        generator.parameters(),
-        lr=args.lr * g_reg_ratio,
-        betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),
-    )
-    d_optim = optim.Adam(
-        discriminator.parameters(),
-        lr=args.lr * d_reg_ratio,
-        betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
-    )
-
-    if args.ckpt is not None:
-        print("load model:", args.ckpt)
-
-        ckpt = torch.load(args.ckpt, map_location=lambda storage, loc: storage)
-
-        try:
-            ckpt_name = os.path.basename(args.ckpt)
-            args.start_iter = int(os.path.splitext(ckpt_name)[0])
-
-        except ValueError:
-            pass
-
-        generator.load_state_dict(ckpt["g"])
-        discriminator.load_state_dict(ckpt["d"])
-        g_ema.load_state_dict(ckpt["g_ema"])
-
-        g_optim.load_state_dict(ckpt["g_optim"])
-        d_optim.load_state_dict(ckpt["d_optim"])
-
-    if args.distributed:
-        generator = nn.parallel.DistributedDataParallel(
-            generator,
-            device_ids=[args.local_rank],
-            output_device=args.local_rank,
-            broadcast_buffers=False,
-        )
-
-        discriminator = nn.parallel.DistributedDataParallel(
-            discriminator,
-            device_ids=[args.local_rank],
-            output_device=args.local_rank,
-            broadcast_buffers=False,
-        )
-
-    transform = transforms.Compose(
-        [
-            transforms.RandomHorizontalFlip(),
-            transforms.ToTensor(),
-            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
-        ]
-    )
-
-    dataset = MultiResolutionDataset(args.path, transform, args.size)
-    loader = data.DataLoader(
-        dataset,
-        batch_size=args.batch,
-        sampler=data_sampler(dataset, shuffle=True, distributed=args.distributed),
-        drop_last=True,
-    )
-
-    if get_rank() == 0 and wandb is not None and args.wandb:
-        wandb.init(project="stylegan 2")
-
-    train(args, loader, generator, discriminator, g_optim, d_optim, g_ema, device)