converter.py

"""Utilities for converting TFHub BigGAN generator weights to PyTorch.

Recommended usage:

To convert all BigGAN variants and generate test samples, use:

```bash
CUDA_VISIBLE_DEVICES=0 python converter.py --generate_samples
```

See `parse_args` for additional options.
"""

import argparse
import os
import sys

import h5py
import torch
import torch.nn as nn
from torchvision.utils import save_image
import tensorflow as tf
import tensorflow_hub as hub
import parse

# import reference biggan from this folder
import biggan_v1 as biggan_for_conversion

# Import model from main folder
sys.path.append('..')
import BigGAN




DEVICE = 'cuda'
HDF5_TMPL = 'biggan-{}.h5'
PTH_TMPL = 'biggan-{}.pth'
MODULE_PATH_TMPL = 'https://tfhub.dev/deepmind/biggan-{}/2'
Z_DIMS = {
  128: 120,
  256: 140,
  512: 128}
RESOLUTIONS = list(Z_DIMS)


def dump_tfhub_to_hdf5(module_path, hdf5_path, redownload=False):
  """Loads TFHub weights and saves them to intermediate HDF5 file.

  Args:
    module_path ([Path-like]): Path to TFHub module.
    hdf5_path ([Path-like]): Path to output HDF5 file.

  Returns:
    [h5py.File]: Loaded hdf5 file containing module weights.
  """
  if os.path.exists(hdf5_path) and (not redownload):
    print('Loading BigGAN hdf5 file from:', hdf5_path)
    return h5py.File(hdf5_path, 'r')

  print('Loading BigGAN module from:', module_path)
  tf.reset_default_graph()
  hub.Module(module_path)
  print('Loaded BigGAN module from:', module_path)

  initializer = tf.global_variables_initializer()
  sess = tf.Session()
  sess.run(initializer)

  print('Saving BigGAN weights to :', hdf5_path)
  h5f = h5py.File(hdf5_path, 'w')
  for var in tf.global_variables():
    val = sess.run(var)
    h5f.create_dataset(var.name, data=val)
    print(f'Saving {var.name} with shape {val.shape}')
  h5f.close()
  return h5py.File(hdf5_path, 'r')


class TFHub2Pytorch(object):

  TF_ROOT = 'module'

  NUM_GBLOCK = {
    128: 5,
    256: 6,
    512: 7
  }

  w = 'w'
  b = 'b'
  u = 'u0'
  v = 'u1'
  gamma = 'gamma'
  beta = 'beta'

  def __init__(self, state_dict, tf_weights, resolution=256, load_ema=True, verbose=False):
    self.state_dict = state_dict
    self.tf_weights = tf_weights
    self.resolution = resolution
    self.verbose = verbose
    if load_ema:
      for name in ['w', 'b', 'gamma', 'beta']:
        setattr(self, name, getattr(self, name) + '/ema_b999900')

  def load(self):
    self.load_generator()
    return self.state_dict

  def load_generator(self):
    GENERATOR_ROOT = os.path.join(self.TF_ROOT, 'Generator')

    for i in range(self.NUM_GBLOCK[self.resolution]):
      name_tf = os.path.join(GENERATOR_ROOT, 'GBlock')
      name_tf += f'_{i}' if i != 0 else ''
      self.load_GBlock(f'GBlock.{i}.', name_tf)

    self.load_attention('attention.', os.path.join(GENERATOR_ROOT, 'attention'))
    self.load_linear('linear', os.path.join(self.TF_ROOT, 'linear'), bias=False)
    self.load_snlinear('G_linear', os.path.join(GENERATOR_ROOT, 'G_Z', 'G_linear'))
    self.load_colorize('colorize', os.path.join(GENERATOR_ROOT, 'conv_2d'))
    self.load_ScaledCrossReplicaBNs('ScaledCrossReplicaBN',
                    os.path.join(GENERATOR_ROOT, 'ScaledCrossReplicaBN'))

  def load_linear(self, name_pth, name_tf, bias=True):
    self.state_dict[name_pth + '.weight'] = self.load_tf_tensor(name_tf, self.w).permute(1, 0)
    if bias:
      self.state_dict[name_pth + '.bias'] = self.load_tf_tensor(name_tf, self.b)

  def load_snlinear(self, name_pth, name_tf, bias=True):
    self.state_dict[name_pth + '.module.weight_u'] = self.load_tf_tensor(name_tf, self.u).squeeze()
    self.state_dict[name_pth + '.module.weight_v'] = self.load_tf_tensor(name_tf, self.v).squeeze()
    self.state_dict[name_pth + '.module.weight_bar'] = self.load_tf_tensor(name_tf, self.w).permute(1, 0)
    if bias:
      self.state_dict[name_pth + '.module.bias'] = self.load_tf_tensor(name_tf, self.b)

  def load_colorize(self, name_pth, name_tf):
    self.load_snconv(name_pth, name_tf)

  def load_GBlock(self, name_pth, name_tf):
    self.load_convs(name_pth, name_tf)
    self.load_HyperBNs(name_pth, name_tf)

  def load_convs(self, name_pth, name_tf):
    self.load_snconv(name_pth + 'conv0', os.path.join(name_tf, 'conv0'))
    self.load_snconv(name_pth + 'conv1', os.path.join(name_tf, 'conv1'))
    self.load_snconv(name_pth + 'conv_sc', os.path.join(name_tf, 'conv_sc'))

  def load_snconv(self, name_pth, name_tf, bias=True):
    if self.verbose:
      print(f'loading: {name_pth} from {name_tf}')
    self.state_dict[name_pth + '.module.weight_u'] = self.load_tf_tensor(name_tf, self.u).squeeze()
    self.state_dict[name_pth + '.module.weight_v'] = self.load_tf_tensor(name_tf, self.v).squeeze()
    self.state_dict[name_pth + '.module.weight_bar'] = self.load_tf_tensor(name_tf, self.w).permute(3, 2, 0, 1)
    if bias:
      self.state_dict[name_pth + '.module.bias'] = self.load_tf_tensor(name_tf, self.b).squeeze()

  def load_conv(self, name_pth, name_tf, bias=True):

    self.state_dict[name_pth + '.weight_u'] = self.load_tf_tensor(name_tf, self.u).squeeze()
    self.state_dict[name_pth + '.weight_v'] = self.load_tf_tensor(name_tf, self.v).squeeze()
    self.state_dict[name_pth + '.weight_bar'] = self.load_tf_tensor(name_tf, self.w).permute(3, 2, 0, 1)
    if bias:
      self.state_dict[name_pth + '.bias'] = self.load_tf_tensor(name_tf, self.b)

  def load_HyperBNs(self, name_pth, name_tf):
    self.load_HyperBN(name_pth + 'HyperBN', os.path.join(name_tf, 'HyperBN'))
    self.load_HyperBN(name_pth + 'HyperBN_1', os.path.join(name_tf, 'HyperBN_1'))

  def load_ScaledCrossReplicaBNs(self, name_pth, name_tf):
    self.state_dict[name_pth + '.bias'] = self.load_tf_tensor(name_tf, self.beta).squeeze()
    self.state_dict[name_pth + '.weight'] = self.load_tf_tensor(name_tf, self.gamma).squeeze()
    self.state_dict[name_pth + '.running_mean'] = self.load_tf_tensor(name_tf + 'bn', 'accumulated_mean')
    self.state_dict[name_pth + '.running_var'] = self.load_tf_tensor(name_tf + 'bn', 'accumulated_var')
    self.state_dict[name_pth + '.num_batches_tracked'] = torch.tensor(
      self.tf_weights[os.path.join(name_tf + 'bn', 'accumulation_counter:0')][()], dtype=torch.float32)

  def load_HyperBN(self, name_pth, name_tf):
    if self.verbose:
      print(f'loading: {name_pth} from {name_tf}')
    beta = name_pth + '.beta_embed.module'
    gamma = name_pth + '.gamma_embed.module'
    self.state_dict[beta + '.weight_u'] = self.load_tf_tensor(os.path.join(name_tf, 'beta'), self.u).squeeze()
    self.state_dict[gamma + '.weight_u'] = self.load_tf_tensor(os.path.join(name_tf, 'gamma'), self.u).squeeze()
    self.state_dict[beta + '.weight_v'] = self.load_tf_tensor(os.path.join(name_tf, 'beta'), self.v).squeeze()
    self.state_dict[gamma + '.weight_v'] = self.load_tf_tensor(os.path.join(name_tf, 'gamma'), self.v).squeeze()
    self.state_dict[beta + '.weight_bar'] = self.load_tf_tensor(os.path.join(name_tf, 'beta'), self.w).permute(1, 0)
    self.state_dict[gamma +
            '.weight_bar'] = self.load_tf_tensor(os.path.join(name_tf, 'gamma'), self.w).permute(1, 0)

    cr_bn_name = name_tf.replace('HyperBN', 'CrossReplicaBN')
    self.state_dict[name_pth + '.bn.running_mean'] = self.load_tf_tensor(cr_bn_name, 'accumulated_mean')
    self.state_dict[name_pth + '.bn.running_var'] = self.load_tf_tensor(cr_bn_name, 'accumulated_var')
    self.state_dict[name_pth + '.bn.num_batches_tracked'] = torch.tensor(
      self.tf_weights[os.path.join(cr_bn_name, 'accumulation_counter:0')][()], dtype=torch.float32)

  def load_attention(self, name_pth, name_tf):

    self.load_snconv(name_pth + 'theta', os.path.join(name_tf, 'theta'), bias=False)
    self.load_snconv(name_pth + 'phi', os.path.join(name_tf, 'phi'), bias=False)
    self.load_snconv(name_pth + 'g', os.path.join(name_tf, 'g'), bias=False)
    self.load_snconv(name_pth + 'o_conv', os.path.join(name_tf, 'o_conv'), bias=False)
    self.state_dict[name_pth + 'gamma'] = self.load_tf_tensor(name_tf, self.gamma)

  def load_tf_tensor(self, prefix, var, device='0'):
    name = os.path.join(prefix, var) + f':{device}'
    return torch.from_numpy(self.tf_weights[name][:])

# Convert from v1: This function maps 
def convert_from_v1(hub_dict, resolution=128):
  weightname_dict = {'weight_u': 'u0', 'weight_bar': 'weight', 'bias': 'bias'}
  convnum_dict = {'conv0': 'conv1', 'conv1': 'conv2', 'conv_sc': 'conv_sc'}
  attention_blocknum = {128: 3, 256: 4, 512: 3}[resolution]
  hub2me = {'linear.weight': 'shared.weight', # This is actually the shared weight 
          # Linear stuff
          'G_linear.module.weight_bar': 'linear.weight', 
          'G_linear.module.bias': 'linear.bias',
          'G_linear.module.weight_u': 'linear.u0',
          # output layer stuff
          'ScaledCrossReplicaBN.weight': 'output_layer.0.gain', 
          'ScaledCrossReplicaBN.bias': 'output_layer.0.bias',
          'ScaledCrossReplicaBN.running_mean': 'output_layer.0.stored_mean',
          'ScaledCrossReplicaBN.running_var': 'output_layer.0.stored_var',
          'colorize.module.weight_bar': 'output_layer.2.weight', 
          'colorize.module.bias': 'output_layer.2.bias',
          'colorize.module.weight_u':  'output_layer.2.u0',
          # Attention stuff
          'attention.gamma': 'blocks.%d.1.gamma' % attention_blocknum, 
          'attention.theta.module.weight_u': 'blocks.%d.1.theta.u0' % attention_blocknum,
          'attention.theta.module.weight_bar': 'blocks.%d.1.theta.weight' % attention_blocknum, 
          'attention.phi.module.weight_u':  'blocks.%d.1.phi.u0' % attention_blocknum,
          'attention.phi.module.weight_bar': 'blocks.%d.1.phi.weight' % attention_blocknum,
          'attention.g.module.weight_u': 'blocks.%d.1.g.u0' % attention_blocknum,
          'attention.g.module.weight_bar': 'blocks.%d.1.g.weight' % attention_blocknum, 
          'attention.o_conv.module.weight_u': 'blocks.%d.1.o.u0' % attention_blocknum,
          'attention.o_conv.module.weight_bar':'blocks.%d.1.o.weight' % attention_blocknum, 
          }

  # Loop over the hub dict and build the hub2me map
  for name in hub_dict.keys():
    if 'GBlock' in name:
      if 'HyperBN' not in name: # it's a conv
        out = parse.parse('GBlock.{:d}.{}.module.{}',name)
        blocknum, convnum, weightname = out
        if weightname not in weightname_dict:
          continue # else hyperBN in 
        out_name = 'blocks.%d.0.%s.%s' % (blocknum, convnum_dict[convnum], weightname_dict[weightname]) # Increment conv number by 1
      else: # hyperbn not conv
        BNnum = 2 if 'HyperBN_1' in name else 1
        if 'embed' in name:        
          out = parse.parse('GBlock.{:d}.{}.module.{}',name)
          blocknum, gamma_or_beta, weightname = out
          if weightname not in weightname_dict: # Ignore weight_v
            continue
          out_name = 'blocks.%d.0.bn%d.%s.%s' % (blocknum, BNnum, 'gain' if 'gamma' in gamma_or_beta else 'bias', weightname_dict[weightname])
        else:
           out = parse.parse('GBlock.{:d}.{}.bn.{}',name)
           blocknum, dummy, mean_or_var = out
           if 'num_batches_tracked' in mean_or_var:
            continue
           out_name = 'blocks.%d.0.bn%d.%s' % (blocknum, BNnum, 'stored_mean' if 'mean' in mean_or_var else 'stored_var')
      hub2me[name] = out_name


  # Invert the hub2me map
  me2hub = {hub2me[item]: item for item in hub2me}
  new_dict = {}
  dimz_dict = {128: 20, 256: 20, 512:16} 
  for item in me2hub:
    # Swap input dim ordering on batchnorm bois to account for my arbitrary change of ordering when concatenating Ys and Zs  
    if ('bn' in item and 'weight' in item) and ('gain' in item or 'bias' in item) and ('output_layer' not in item):
      new_dict[item] = torch.cat([hub_dict[me2hub[item]][:, -128:], hub_dict[me2hub[item]][:, :dimz_dict[resolution]]], 1)
    # Reshape the first linear weight, bias, and u0
    elif item == 'linear.weight':
      new_dict[item] = hub_dict[me2hub[item]].contiguous().view(4, 4, 96 * 16, -1).permute(2,0,1,3).contiguous().view(-1,dimz_dict[resolution])
    elif item == 'linear.bias':
      new_dict[item] = hub_dict[me2hub[item]].view(4, 4, 96  * 16).permute(2,0,1).contiguous().view(-1)
    elif item == 'linear.u0':
       new_dict[item] = hub_dict[me2hub[item]].view(4, 4, 96  * 16).permute(2,0,1).contiguous().view(1, -1)
    elif me2hub[item] == 'linear.weight': # THIS IS THE SHARED WEIGHT NOT THE FIRST LINEAR LAYER
      # Transpose shared weight so that it's an embedding
      new_dict[item] = hub_dict[me2hub[item]].t()
    elif 'weight_u' in me2hub[item]: # Unsqueeze u0s    
      new_dict[item] = hub_dict[me2hub[item]].unsqueeze(0)
    else:
      new_dict[item] = hub_dict[me2hub[item]]      
  return new_dict

def get_config(resolution):
  attn_dict = {128: '64', 256: '128', 512: '64'}
  dim_z_dict = {128: 120, 256: 140, 512: 128}
  config = {'G_param': 'SN', 'D_param': 'SN', 
           'G_ch': 96, 'D_ch': 96, 
           'D_wide': True, 'G_shared': True, 
           'shared_dim': 128, 'dim_z': dim_z_dict[resolution], 
           'hier': True, 'cross_replica': False, 
           'mybn': False, 'G_activation': nn.ReLU(inplace=True),
           'G_attn': attn_dict[resolution],
           'norm_style': 'bn',
           'G_init': 'ortho', 'skip_init': True, 'no_optim': True,
           'G_fp16': False, 'G_mixed_precision': False,
           'accumulate_stats': False, 'num_standing_accumulations': 16, 
           'G_eval_mode': True,
           'BN_eps': 1e-04, 'SN_eps': 1e-04, 
           'num_G_SVs': 1, 'num_G_SV_itrs': 1, 'resolution': resolution, 
           'n_classes': 1000}
  return config


def convert_biggan(resolution, weight_dir, redownload=False, no_ema=False, verbose=False):
  module_path = MODULE_PATH_TMPL.format(resolution)
  hdf5_path = os.path.join(weight_dir, HDF5_TMPL.format(resolution))
  pth_path = os.path.join(weight_dir, PTH_TMPL.format(resolution))

  tf_weights = dump_tfhub_to_hdf5(module_path, hdf5_path, redownload=redownload)
  G_temp = getattr(biggan_for_conversion, f'Generator{resolution}')()
  state_dict_temp = G_temp.state_dict()

  converter = TFHub2Pytorch(state_dict_temp, tf_weights, resolution=resolution,
                load_ema=(not no_ema), verbose=verbose)
  state_dict_v1 = converter.load()
  state_dict = convert_from_v1(state_dict_v1, resolution)
  # Get the config, build the model
  config = get_config(resolution)
  G = BigGAN.Generator(**config)
  G.load_state_dict(state_dict, strict=False) # Ignore missing sv0 entries
  torch.save(state_dict, pth_path)
  
  # output_location ='pretrained_weights/TFHub-PyTorch-128.pth'
  
  return G


def generate_sample(G, z_dim, batch_size, filename, parallel=False):
  
  G.eval()
  G.to(DEVICE)
  with torch.no_grad():
    z = torch.randn(batch_size, G.dim_z).to(DEVICE)
    y = torch.randint(low=0, high=1000, size=(batch_size,), 
        device=DEVICE, dtype=torch.int64, requires_grad=False)
    if parallel:
      images = nn.parallel.data_parallel(G, (z, G.shared(y)))
    else:
      images = G(z, G.shared(y))
  save_image(images, filename, scale_each=True, normalize=True)

def parse_args():
  usage = 'Parser for conversion script.'
  parser = argparse.ArgumentParser(description=usage)
  parser.add_argument(
    '--resolution', '-r', type=int, default=None, choices=[128, 256, 512],
    help='Resolution of TFHub module to convert. Converts all resolutions if None.')
  parser.add_argument(
    '--redownload', action='store_true', default=False,
    help='Redownload weights and overwrite current hdf5 file, if present.')
  parser.add_argument(
    '--weights_dir', type=str, default='pretrained_weights')
  parser.add_argument(
    '--samples_dir', type=str, default='pretrained_samples')
  parser.add_argument(
    '--no_ema', action='store_true', default=False,
    help='Do not load ema weights.')
  parser.add_argument(
    '--verbose', action='store_true', default=False,
    help='Additionally logging.')
  parser.add_argument(
    '--generate_samples', action='store_true', default=False,
    help='Generate test sample with pretrained model.')
  parser.add_argument(
    '--batch_size', type=int, default=64,
    help='Batch size used for test sample.')
  parser.add_argument(
    '--parallel', action='store_true', default=False,
    help='Parallelize G?')     
  args = parser.parse_args()
  return args


if __name__ == '__main__':

  args = parse_args()
  os.makedirs(args.weights_dir, exist_ok=True)
  os.makedirs(args.samples_dir, exist_ok=True)

  if args.resolution is not None:
    G = convert_biggan(args.resolution, args.weights_dir,
               redownload=args.redownload,
               no_ema=args.no_ema, verbose=args.verbose)
    if args.generate_samples:
      filename = os.path.join(args.samples_dir, f'biggan{args.resolution}_samples.jpg')
      print('Generating samples...')
      generate_sample(G, Z_DIMS[args.resolution], args.batch_size, filename, args.parallel)
  else:
    for res in RESOLUTIONS:
      G = convert_biggan(res, args.weights_dir,
                 redownload=args.redownload,
                 no_ema=args.no_ema, verbose=args.verbose)
      if args.generate_samples:
        filename = os.path.join(args.samples_dir, f'biggan{res}_samples.jpg')
        print('Generating samples...')
        generate_sample(G, Z_DIMS[res], args.batch_size, filename, args.parallel)