infer_wavenet.py 文件源码

def generate_and_save_samples(sample_fn, length, count, dir, rate, levels):
    def save_samples(data):
        data = (data * np.reshape(np.arange(levels) / (levels-1), [levels, 1, 1])).sum(
            axis=1, keepdims=True)
        value = np.iinfo(np.int16).max
        audio = (utils.inverse_mulaw(data * 2 - 1) * value).astype(np.int16)
        for idx, sample in enumerate(audio):
            filename = os.path.join(dir, 'sample_{}.wav'.format(idx))
            wavfile.write(filename, rate, np.squeeze(sample))

    samples = chainer.Variable(
        chainer.cuda.cupy.zeros([count, levels, 1, length], dtype='float32'))
    one_hot_ref = chainer.cuda.cupy.eye(levels).astype('float32')

    with tqdm.tqdm(total=length) as bar:
        for i in range(length):
            probs = F.softmax(sample_fn(samples))[:, :, 0, 0, i]
            samples.data[:, :, 0, i] = one_hot_ref[utils.sample_from(probs.data.get())]
            bar.update()

    samples.to_cpu()
    save_samples(samples.data)