profiler.py

import argparse
import time

import numpy
import chainer

from WaveGlow import Glow
import params

parser = argparse.ArgumentParser()
parser.add_argument('--gpu', '-g', type=int, default=-1,
                    help='GPU ID (negative value indicates CPU)')
parser.add_argument('--length', '-l', type=int, default=1,
                    help='length to generate(sec)')

args = parser.parse_args()
if args.gpu != [-1]:
    chainer.cuda.set_max_workspace_size(2 * 512 * 1024 * 1024)
    chainer.global_config.autotune = True

# make model
glow = Glow(
    params.hop_length, params.n_mels, 1,
    params.squeeze_factor, params.n_flows, params.n_layers,
    params.wn_channel, params.early_every, params.early_size,
    params.var)

# make input
condition = numpy.empty(
    (1, params.n_mels, args.length * params.sr // params.hop_length), numpy.float32)

if args.gpu >= 0:
    use_gpu = True
    chainer.cuda.get_device_from_id(args.gpu).use()
else:
    use_gpu = False

# forward
if use_gpu:
    condition = chainer.cuda.to_gpu(condition, device=args.gpu)
    glow.to_gpu(device=args.gpu)
condition = chainer.Variable(condition)

with chainer.using_config('enable_backprop', False):
    output = glow.generate(condition)
print('Generate', output.shape[-1] / params.sr, 'secs...')

times = 10
with chainer.using_config('enable_backprop', False):
    for i in range(times):
        start = time.time()
        glow.generate(condition)
        end = time.time()
        print('It took', (end - start) / times, 'secs to generate.')