add export script

convert_to_pb.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+# Copyright    2023  Xiaomi Corp.        (authors: Fangjun Kuang)
+
+# Please see ./run.sh for usages
+import argparse
+import os
+
+import tensorflow as tf
+
+
+# Code in the following function is modified from
+# https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc
+def freeze_graph(model_dir, output_node_names, output_filename):
+    """Extract the sub graph defined by the output nodes and convert all its
+    variables into constant
+
+    Args:
+      model_dir:
+        the root folder containing the checkpoint state file
+      output_node_names:
+        a string, containing all the output node's names, comma separated
+      output_filename:
+        Filename to save the graph.
+    """
+    if not tf.compat.v1.gfile.Exists(model_dir):
+        raise AssertionError(
+            "Export directory doesn't exists. Please specify an export "
+            "directory: %s" % model_dir
+        )
+
+    if not output_node_names:
+        print("You need to supply the name of a node to --output_node_names.")
+        return -1
+
+    # We retrieve our checkpoint fullpath
+    checkpoint = tf.train.get_checkpoint_state(model_dir)
+    input_checkpoint = checkpoint.model_checkpoint_path
+
+    # We precise the file fullname of our freezed graph
+    absolute_model_dir = "/".join(input_checkpoint.split("/")[:-1])
+    output_graph = output_filename
+
+    # We clear devices to allow TensorFlow to control on which device it will load operations
+    clear_devices = True
+
+    # We start a session using a temporary fresh Graph
+    with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+        # We import the meta graph in the current default Graph
+        saver = tf.compat.v1.train.import_meta_graph(
+            input_checkpoint + ".meta", clear_devices=clear_devices
+        )
+
+        # We restore the weights
+        saver.restore(sess, input_checkpoint)
+
+        # We use a built-in TF helper to export variables to constants
+        output_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
+            sess,  # The session is used to retrieve the weights
+            tf.compat.v1.get_default_graph().as_graph_def(),  # The graph_def is used to retrieve the nodes
+            output_node_names.split(
+                ","
+            ),  # The output node names are used to select the usefull nodes
+        )
+
+        # Finally we serialize and dump the output graph to the filesystem
+        with tf.compat.v1.gfile.GFile(output_graph, "wb") as f:
+            f.write(output_graph_def.SerializeToString())
+        print("%d ops in the final graph." % len(output_graph_def.node))
+
+    return output_graph_def
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model-dir", type=str, default="", help="Model folder to export"
+    )
+    parser.add_argument(
+        "--output-node-names",
+        type=str,
+        default="vocals_spectrogram/mul,accompaniment_spectrogram/mul",
+        help="The name of the output nodes, comma separated.",
+    )
+
+    parser.add_argument(
+        "--output-filename",
+        type=str,
+    )
+    args = parser.parse_args()
+
+    freeze_graph(args.model_dir, args.output_node_names, args.output_filename)

convert_to_torch.py

@@ -0,0 +1,240 @@
+#!/usr/bin/env python3
+# Copyright    2023  Xiaomi Corp.        (authors: Fangjun Kuang)
+
+# Please see ./run.sh for usage
+
+import argparse
+
+import numpy as np
+import tensorflow as tf
+import torch
+import torch.nn as nn
+from unet import UNet
+
+
+def load_graph(frozen_graph_filename):
+    # This function is modified from
+    # https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc
+
+    # We load the protobuf file from the disk and parse it to retrieve the
+    # unserialized graph_def
+    with tf.compat.v1.gfile.GFile(frozen_graph_filename, "rb") as f:
+        graph_def = tf.compat.v1.GraphDef()
+        graph_def.ParseFromString(f.read())
+
+    # Then, we import the graph_def into a new Graph and returns it
+    with tf.Graph().as_default() as graph:
+        # The name var will prefix every op/nodes in your graph
+        # Since we load everything in a new graph, this is not needed
+        #  tf.import_graph_def(graph_def, name="prefix")
+        tf.import_graph_def(graph_def, name="")
+    return graph
+
+
+def generate_waveform():
+    np.random.seed(20230821)
+    waveform = np.random.rand(60 * 44100).astype(np.float32)
+
+    # (num_samples, num_channels)
+    waveform = waveform.reshape(-1, 2)
+    return waveform
+
+
+def get_param(graph, name):
+    with tf.compat.v1.Session(graph=graph) as sess:
+        constant_ops = [op for op in sess.graph.get_operations() if op.type == "Const"]
+        for constant_op in constant_ops:
+            if constant_op.name != name:
+                continue
+
+            value = sess.run(constant_op.outputs[0])
+            return torch.from_numpy(value)
+
+
+@torch.no_grad()
+def main(name):
+    graph = load_graph(f"./2stems/frozen_{name}_model.pb")
+    #  for op in graph.get_operations():
+    #      print(op.name)
+    x = graph.get_tensor_by_name("waveform:0")
+    #  y = graph.get_tensor_by_name("Reshape:0")
+    y0 = graph.get_tensor_by_name("strided_slice_3:0")
+    #  y1 = graph.get_tensor_by_name("leaky_re_lu_5/LeakyRelu:0")
+    #  y1 = graph.get_tensor_by_name("conv2d_5/BiasAdd:0")
+    #  y1 = graph.get_tensor_by_name("conv2d_transpose/BiasAdd:0")
+    #  y1 = graph.get_tensor_by_name("re_lu/Relu:0")
+    #  y1 = graph.get_tensor_by_name("batch_normalization_6/cond/FusedBatchNorm_1:0")
+    #  y1 = graph.get_tensor_by_name("concatenate/concat:0")
+    #  y1 = graph.get_tensor_by_name("concatenate_1/concat:0")
+    #  y1 = graph.get_tensor_by_name("concatenate_4/concat:0")
+    #  y1 = graph.get_tensor_by_name("batch_normalization_11/cond/FusedBatchNorm_1:0")
+    #  y1 = graph.get_tensor_by_name("conv2d_6/Sigmoid:0")
+    y1 = graph.get_tensor_by_name(f"{name}_spectrogram/mul:0")
+
+    unet = UNet()
+    unet.eval()
+
+    # For the conv2d in tensorflow, weight shape is (kernel_h, kernel_w, in_channel, out_channel)
+    # default input shape is NHWC
+
+    # For the conv2d in torch, weight shape is (out_channel, in_channel, kernel_h, kernel_w)
+    # default input shape is NCHW
+    state_dict = unet.state_dict()
+    #  print(list(state_dict.keys()))
+
+    if name == "vocals":
+        state_dict["conv.weight"] = get_param(graph, "conv2d/kernel").permute(
+            3, 2, 0, 1
+        )
+        state_dict["conv.bias"] = get_param(graph, "conv2d/bias")
+
+        state_dict["bn.weight"] = get_param(graph, "batch_normalization/gamma")
+        state_dict["bn.bias"] = get_param(graph, "batch_normalization/beta")
+        state_dict["bn.running_mean"] = get_param(
+            graph, "batch_normalization/moving_mean"
+        )
+        state_dict["bn.running_var"] = get_param(
+            graph, "batch_normalization/moving_variance"
+        )
+
+        conv_offset = 0
+        bn_offset = 0
+    else:
+        state_dict["conv.weight"] = get_param(graph, "conv2d_7/kernel").permute(
+            3, 2, 0, 1
+        )
+        state_dict["conv.bias"] = get_param(graph, "conv2d_7/bias")
+
+        state_dict["bn.weight"] = get_param(graph, "batch_normalization_12/gamma")
+        state_dict["bn.bias"] = get_param(graph, "batch_normalization_12/beta")
+        state_dict["bn.running_mean"] = get_param(
+            graph, "batch_normalization_12/moving_mean"
+        )
+        state_dict["bn.running_var"] = get_param(
+            graph, "batch_normalization_12/moving_variance"
+        )
+        conv_offset = 7
+        bn_offset = 12
+
+    for i in range(1, 6):
+        state_dict[f"conv{i}.weight"] = get_param(
+            graph, f"conv2d_{i+conv_offset}/kernel"
+        ).permute(3, 2, 0, 1)
+        state_dict[f"conv{i}.bias"] = get_param(graph, f"conv2d_{i+conv_offset}/bias")
+        if i >= 5:
+            continue
+        state_dict[f"bn{i}.weight"] = get_param(
+            graph, f"batch_normalization_{i+bn_offset}/gamma"
+        )
+        state_dict[f"bn{i}.bias"] = get_param(
+            graph, f"batch_normalization_{i+bn_offset}/beta"
+        )
+        state_dict[f"bn{i}.running_mean"] = get_param(
+            graph, f"batch_normalization_{i+bn_offset}/moving_mean"
+        )
+        state_dict[f"bn{i}.running_var"] = get_param(
+            graph, f"batch_normalization_{i+bn_offset}/moving_variance"
+        )
+
+    if name == "vocals":
+        state_dict["up1.weight"] = get_param(graph, "conv2d_transpose/kernel").permute(
+            3, 2, 0, 1
+        )
+        state_dict["up1.bias"] = get_param(graph, "conv2d_transpose/bias")
+
+        state_dict["bn5.weight"] = get_param(graph, "batch_normalization_6/gamma")
+        state_dict["bn5.bias"] = get_param(graph, "batch_normalization_6/beta")
+        state_dict["bn5.running_mean"] = get_param(
+            graph, "batch_normalization_6/moving_mean"
+        )
+        state_dict["bn5.running_var"] = get_param(
+            graph, "batch_normalization_6/moving_variance"
+        )
+        conv_offset = 0
+        bn_offset = 0
+    else:
+        state_dict["up1.weight"] = get_param(
+            graph, "conv2d_transpose_6/kernel"
+        ).permute(3, 2, 0, 1)
+        state_dict["up1.bias"] = get_param(graph, "conv2d_transpose_6/bias")
+
+        state_dict["bn5.weight"] = get_param(graph, "batch_normalization_18/gamma")
+        state_dict["bn5.bias"] = get_param(graph, "batch_normalization_18/beta")
+        state_dict["bn5.running_mean"] = get_param(
+            graph, "batch_normalization_18/moving_mean"
+        )
+        state_dict["bn5.running_var"] = get_param(
+            graph, "batch_normalization_18/moving_variance"
+        )
+        conv_offset = 6
+        bn_offset = 12
+
+    for i in range(1, 6):
+        state_dict[f"up{i+1}.weight"] = get_param(
+            graph, f"conv2d_transpose_{i+conv_offset}/kernel"
+        ).permute(3, 2, 0, 1)
+
+        state_dict[f"up{i+1}.bias"] = get_param(
+            graph, f"conv2d_transpose_{i+conv_offset}/bias"
+        )
+
+        state_dict[f"bn{5+i}.weight"] = get_param(
+            graph, f"batch_normalization_{6+i+bn_offset}/gamma"
+        )
+        state_dict[f"bn{5+i}.bias"] = get_param(
+            graph, f"batch_normalization_{6+i+bn_offset}/beta"
+        )
+        state_dict[f"bn{5+i}.running_mean"] = get_param(
+            graph, f"batch_normalization_{6+i+bn_offset}/moving_mean"
+        )
+        state_dict[f"bn{5+i}.running_var"] = get_param(
+            graph, f"batch_normalization_{6+i+bn_offset}/moving_variance"
+        )
+
+    if name == "vocals":
+        state_dict["up7.weight"] = get_param(graph, "conv2d_6/kernel").permute(
+            3, 2, 0, 1
+        )
+        state_dict["up7.bias"] = get_param(graph, "conv2d_6/bias")
+    else:
+        state_dict["up7.weight"] = get_param(graph, "conv2d_13/kernel").permute(
+            3, 2, 0, 1
+        )
+        state_dict["up7.bias"] = get_param(graph, "conv2d_13/bias")
+
+    unet.load_state_dict(state_dict)
+
+    with tf.compat.v1.Session(graph=graph) as sess:
+        y0_out, y1_out = sess.run([y0, y1], feed_dict={x: generate_waveform()})
+        #  y0_out = sess.run(y0, feed_dict={x: generate_waveform()})
+        #  y1_out = sess.run(y1, feed_dict={x: generate_waveform()})
+        #  print(y0_out.shape)
+        #  print(y1_out.shape)
+
+    # for the batchnormalization in tensorflow,
+    # default input shape is NHWC
+
+    # for the batchnormalization in torch,
+    # default input shape is NCHW
+
+    # NHWC to NCHW
+    torch_y1_out = unet(torch.from_numpy(y0_out).permute(0, 3, 1, 2))
+
+    #  print(torch_y1_out.shape, torch.from_numpy(y1_out).permute(0, 3, 1, 2).shape)
+    assert torch.allclose(
+        torch_y1_out, torch.from_numpy(y1_out).permute(0, 3, 1, 2), atol=1e-1
+    ), ((torch_y1_out - torch.from_numpy(y1_out).permute(0, 3, 1, 2)).abs().max())
+    torch.save(unet.state_dict(), f"2stems/{name}.pt")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--name",
+        type=str,
+        required=True,
+        choices=["vocals", "accompaniment"],
+    )
+    args = parser.parse_args()
+    print(vars(args))
+    main(args.name)

run.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# Copyright    2023  Xiaomi Corp.        (authors: Fangjun Kuang)
+
+if [ ! -f 2stems.tar.gz ]; then
+  wget https://github.com/deezer/spleeter/releases/download/v1.4.0/2stems.tar.gz
+fi
+
+if [ ! -d ./2stems ]; then
+  mkdir -p 2stems
+  cd 2stems
+  tar xvf ../2stems.tar.gz
+  cd ..
+fi
+
+if [ ! -f 2stems/frozen_vocals_model.pb ]; then
+  python3 ./convert_to_pb.py \
+    --model-dir ./2stems \
+    --output-node-names vocals_spectrogram/mul \
+    --output-filename ./2stems/frozen_vocals_model.pb
+fi
+
+if [ ! -f 2stems/frozen_accompaniment_model.pb ]; then
+  python3 ./convert_to_pb.py \
+    --model-dir ./2stems \
+    --output-node-names accompaniment_spectrogram/mul \
+    --output-filename ./2stems/frozen_accompaniment_model.pb
+fi
+
+python3 ./convert_to_torch.py --name vocals
+python3 ./convert_to_torch.py --name accompaniment

unet.py

@@ -0,0 +1,150 @@
+# Copyright    2023  Xiaomi Corp.        (authors: Fangjun Kuang)
+
+import torch
+
+
+class UNet(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(2, 16, kernel_size=5, stride=(2, 2), padding=0)
+        self.bn = torch.nn.BatchNorm2d(
+            16, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+        #
+        self.conv1 = torch.nn.Conv2d(16, 32, kernel_size=5, stride=(2, 2), padding=0)
+        self.bn1 = torch.nn.BatchNorm2d(
+            32, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=5, stride=(2, 2), padding=0)
+        self.bn2 = torch.nn.BatchNorm2d(
+            64, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.conv3 = torch.nn.Conv2d(64, 128, kernel_size=5, stride=(2, 2), padding=0)
+        self.bn3 = torch.nn.BatchNorm2d(
+            128, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.conv4 = torch.nn.Conv2d(128, 256, kernel_size=5, stride=(2, 2), padding=0)
+        self.bn4 = torch.nn.BatchNorm2d(
+            256, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.conv5 = torch.nn.Conv2d(256, 512, kernel_size=5, stride=(2, 2), padding=0)
+
+        self.up1 = torch.nn.ConvTranspose2d(512, 256, kernel_size=5, stride=2)
+        self.bn5 = torch.nn.BatchNorm2d(
+            256, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.up2 = torch.nn.ConvTranspose2d(512, 128, kernel_size=5, stride=2)
+        self.bn6 = torch.nn.BatchNorm2d(
+            128, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.up3 = torch.nn.ConvTranspose2d(256, 64, kernel_size=5, stride=2)
+        self.bn7 = torch.nn.BatchNorm2d(
+            64, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.up4 = torch.nn.ConvTranspose2d(128, 32, kernel_size=5, stride=2)
+        self.bn8 = torch.nn.BatchNorm2d(
+            32, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.up5 = torch.nn.ConvTranspose2d(64, 16, kernel_size=5, stride=2)
+        self.bn9 = torch.nn.BatchNorm2d(
+            16, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        self.up6 = torch.nn.ConvTranspose2d(32, 1, kernel_size=5, stride=2)
+        self.bn10 = torch.nn.BatchNorm2d(
+            1, track_running_stats=True, eps=1e-3, momentum=0.01
+        )
+
+        # output logit is False, so we need self.up7
+        self.up7 = torch.nn.Conv2d(1, 2, kernel_size=4, dilation=2, padding=3)
+
+    def forward(self, x):
+        in_x = x
+        # in_x is (3, 2, 512, 1024) = (T, 2, 512, 1024)
+        x = torch.nn.functional.pad(x, (1, 2, 1, 2), "constant", 0)
+        conv1 = self.conv(x)
+        batch1 = self.bn(conv1)
+        rel1 = torch.nn.functional.leaky_relu(batch1, negative_slope=0.2)
+
+        x = torch.nn.functional.pad(rel1, (1, 2, 1, 2), "constant", 0)
+        conv2 = self.conv1(x)  # (3, 32, 128, 256)
+        batch2 = self.bn1(conv2)
+        rel2 = torch.nn.functional.leaky_relu(
+            batch2, negative_slope=0.2
+        )  # (3, 32, 128, 256)
+
+        x = torch.nn.functional.pad(rel2, (1, 2, 1, 2), "constant", 0)
+        conv3 = self.conv2(x)  # (3, 64, 64, 128)
+        batch3 = self.bn2(conv3)
+        rel3 = torch.nn.functional.leaky_relu(
+            batch3, negative_slope=0.2
+        )  # (3, 64, 64, 128)
+
+        x = torch.nn.functional.pad(rel3, (1, 2, 1, 2), "constant", 0)
+        conv4 = self.conv3(x)  # (3, 128, 32, 64)
+        batch4 = self.bn3(conv4)
+        rel4 = torch.nn.functional.leaky_relu(
+            batch4, negative_slope=0.2
+        )  # (3, 128, 32, 64)
+
+        x = torch.nn.functional.pad(rel4, (1, 2, 1, 2), "constant", 0)
+        conv5 = self.conv4(x)  # (3, 256, 16, 32)
+        batch5 = self.bn4(conv5)
+        rel6 = torch.nn.functional.leaky_relu(
+            batch5, negative_slope=0.2
+        )  # (3, 256, 16, 32)
+
+        x = torch.nn.functional.pad(rel6, (1, 2, 1, 2), "constant", 0)
+        conv6 = self.conv5(x)  # (3, 512, 8, 16)
+
+        up1 = self.up1(conv6)
+        up1 = up1[:, :, 1:-2, 1:-2]  # (3, 256, 16, 32)
+        up1 = torch.nn.functional.relu(up1)
+        batch7 = self.bn5(up1)
+        merge1 = torch.cat([conv5, batch7], axis=1)  # (3, 512, 16, 32)
+
+        up2 = self.up2(merge1)
+        up2 = up2[:, :, 1:-2, 1:-2]
+        up2 = torch.nn.functional.relu(up2)
+        batch8 = self.bn6(up2)
+
+        merge2 = torch.cat([conv4, batch8], axis=1)  # (3, 256, 32, 64)
+
+        up3 = self.up3(merge2)
+        up3 = up3[:, :, 1:-2, 1:-2]
+        up3 = torch.nn.functional.relu(up3)
+        batch9 = self.bn7(up3)
+
+        merge3 = torch.cat([conv3, batch9], axis=1)  # (3, 128, 64, 128)
+
+        up4 = self.up4(merge3)
+        up4 = up4[:, :, 1:-2, 1:-2]
+        up4 = torch.nn.functional.relu(up4)
+        batch10 = self.bn8(up4)
+
+        merge4 = torch.cat([conv2, batch10], axis=1)  # (3, 64, 128, 256)
+
+        up5 = self.up5(merge4)
+        up5 = up5[:, :, 1:-2, 1:-2]
+        up5 = torch.nn.functional.relu(up5)
+        batch11 = self.bn9(up5)
+
+        merge5 = torch.cat([conv1, batch11], axis=1)  # (3, 32, 256, 512)
+
+        up6 = self.up6(merge5)
+        up6 = up6[:, :, 1:-2, 1:-2]
+        up6 = torch.nn.functional.relu(up6)
+        batch12 = self.bn10(up6)  # (3, 1, 512, 1024)  = (T, 1, 512, 1024)
+
+        up7 = self.up7(batch12)
+        up7 = torch.sigmoid(up7)  # (3, 2, 512, 1024)
+
+        return up7 * in_x