wavenet.py

def mu_law(x, mu):
    return np.sign(x) * np.log(1 + mu * np.abs(x)) / np.log(1 + mu)

# value shape is [batches, width, D]
# filters shape is [width, D, C]
def causal_atrous_conv1d(value, filters, rate, padding):
    # Using height in 2-D as the 1-D.
    value_2d = tf.expand_dims(value, 2)
    filters_2d = tf.expand_dims(filters, 1)
    # Note that for filters using 'SAME' padding, padding zeros are added to the end of the input.
    # This means that for causal convolutions, we must shift the output right.
    # add zeros to the start and remove the future values from the end.
    
    atr_conv_1d = tf.squeeze(tf.nn.atrous_conv2d(value_2d, filters_2d, rate, padding), [2])
    # atr_conv_1d shape is [batches, width, C]
    
    filter_shape = array_ops.shape(filters)
    filter_width = filter_shape[0]
    filter_width_up = filter_width + (filter_width - 1) * (rate - 1)
    pad_width = filter_width_up - 1
    pad_left = pad_width // 2
    pad_right = pad_width - pad_left
    # We want to shift the result so that acausal values are removed.
    # Any value in the output that makes use of right padding values are acausal.
    # So, we remove pad_right elements from the end, and add as many zeros to the beginning.
    C = array_ops.shape(atr_conv_1d)[2]
    causal = tf.pad(tf.slice(atr_conv_1d, [0, 0, 0], [batches, filter_width - pad_right, C]),
                    [[0, 0], [pad_right, 0], [0, 0]])
    return causal

# Returns a tuple of output to the next layer and skip output.
# The shape of x is [batches, width, D]
# w1, w2 shapes are [width, D, C]
# cw shape is [width, C, B]
def gated_unit(x, w1, w2, cw, dilation):
    # x shape is [batches, width, D]
    dilated1 = causal_atrous_conv1d(x, w1, dilation, 'SAME')
    dilated2 = causal_atrous_conv1d(x, w2, dilation, 'SAME')
    z = tf.mul(tf.tanh(dilated1), tf.sigmoid(dilated2))
    # dilated1, dilated2, z shapes are [batches, width, C]
    combined = tf.nn.conv1d(z, cw, 1, 'SAME')
    output = combined + x
    # combined and output shapes are [batches, width, B]
    return (output, combined)

# Returns a tuple of (output, non-softmaxed-logits output)
# The non-softmaxed output is used for the loss calculation.
# D = 1
# The shape of x is [batches, width, D]
# The shape of output is [width, 256]
# w1, w2 shapes are [width, 1, C]
# cw shape is [width, C, B]
# co1 shape is [width, B, A]
# co2 shape is [width, A, 256]
# Dilations is an array of [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1, 2, ..., 512]
def layers(x, dilations, w1, w2, cw, co1, co2):
    next_input = x
    skip_connections = []
    for (dilation, w1, w2,cw) in zip(dilations, w1, w2):
        (output, skip) = gated_unit(next_input, w1)
        # output and skip shapes are [batches, width, B]
        next_input = output
        skip_connections.append(skip)
    sum_of_skip_connections = reduce(lambda l, r: l + r, skip_connections)
    # sum_of_skip_connections shape is [batches, width, B]
    relu1 = tf.nn.relu(sum_of_skip_connections)
    relu2 = tf.nn.relu(tf.nn.conv1d(relu1, co1, 1, 'SAME'))
    raw_output = tf.nn.conv1d(relu2, co2, 1, 'SAME')
    output = tf.nn.softmax(raw_output)
    return (output, raw_output)

def create(parameters):
    x = tf.placeholder(tf.float32, shape=(None, parameters['width']), name='input')
    mu_law_x = mu_law(x, 255)
    # ... Initialize and define the variables here...
    model= {
        x: x
    }
    return model