Spaces:
Sleeping
Sleeping
| # Copyright 2020 Erik Härkönen. All rights reserved. | |
| # This file is licensed to you 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 REPRESENTATIONS | |
| # OF ANY KIND, either express or implied. See the License for the specific language | |
| # governing permissions and limitations under the License. | |
| import torch, numpy as np | |
| from types import SimpleNamespace | |
| import itertools | |
| import sys | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).parent.parent)) | |
| from models import get_model | |
| from config import Config | |
| torch.backends.cudnn.benchmark = True | |
| has_gpu = torch.cuda.is_available() | |
| device = torch.device('cuda' if has_gpu else 'cpu') | |
| B = 2 # test batch support | |
| models = [ | |
| ('BigGAN-128', 'husky'), | |
| ('BigGAN-256', 'husky'), | |
| ('BigGAN-512', 'husky'), | |
| ('StyleGAN', 'ffhq'), | |
| ('StyleGAN2', 'ffhq'), | |
| ] | |
| for model_name, classname in models: | |
| with torch.no_grad(): | |
| model = get_model(model_name, classname, device).to(device) | |
| print(f'Testing {model_name}-{classname}', end='') | |
| n_latents = model.get_max_latents() | |
| assert n_latents > 1, 'Model reports max_latents=1' | |
| #if hasattr(model, 'use_w'): | |
| # model.use_w() | |
| seed = 1234 | |
| torch.manual_seed(seed) | |
| np.random.seed(seed) | |
| latents = [model.sample_latent(B, seed=seed) for _ in range(10)] | |
| # Test that partial-forward supports layerwise latent inputs | |
| try: | |
| layer_name, _ = list(model.named_modules())[-1] | |
| _ = model.partial_forward(n_latents*[latents[0]], layer_name) | |
| except Exception as e: | |
| print('Error:', e) | |
| raise RuntimeError(f"{model_name} partial forward doesn't support layerwise latent!") | |
| # Test that layerwise and single give same result | |
| for z in latents: | |
| torch.manual_seed(0) | |
| np.random.seed(0) | |
| out1 = model.forward(z) | |
| torch.manual_seed(0) | |
| np.random.seed(0) | |
| out2 = model.forward(n_latents*[z]) | |
| dist_rel = (torch.abs(out1 - out2).sum() / out1.sum()).item() | |
| assert dist_rel < 1e-3, f'Layerwise latent mode working incorrectly for model {model_name}-{classname}: difference = {dist_rel*100}%' | |
| print('.', end='') | |
| print('OK!') | |