feat: differentiable PyTorch backend #31
Conversation
|
@yoyololicon thats great 🙌 Have you seen our implementation in https://github.com/sigsep/open-unmix-pytorch/blob/master/openunmix/filtering.py ?
|
Thanks for the info, I haven't checked this one before.
Did you mean using PyTorch native complex type? Yes.
Sure. I think we can also profile the memory usage as well.
As I can remember, we haven't directly compared them before. I agree we should add a regression test and report the numbers. |
|
Thanks a lot, that's great !! I remember talking to you about doing a PR for norbert, that's fantastic you took time to do it. If you could indeed do the few tests @faroit mentions, that would be perfect best |
@aliutkus Yeah, I remember it, too. 😄
Theoretically, it can backprop any number of iterations, but we always set it to one iteration in our training. |
|
@faroit @aliutkus I added the regression tests and should be ready for review. ProfileNorbert TorchOpen-UnmixBenchmarkTesting different combinations of sources, iterations, and threads, with other parameters fixed to a reasonable value. The script I used to produce those experimentsimport torch
from torch.profiler import profile, record_function, ProfilerActivity
import torch.utils.benchmark as benchmark
from itertools import product
import norbert.torch as norbert
from openunmix import filtering
nb_frames = 100
nb_bins = 513
nb_channels = 2
nb_sources = [4, 8]
nb_iterations = [0, 1, 3]
x = torch.randn(nb_frames, nb_bins, nb_channels) + 1j * \
torch.randn(nb_frames, nb_bins, nb_channels)
x_as_real = torch.view_as_real(x)
v = torch.rand(nb_frames, nb_bins, nb_channels, 4)
with profile(activities=[ProfilerActivity.CPU], profile_memory=True) as prof:
norbert.wiener(v[None, ...], x[None, ...])
print("Norbert Torch")
print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=5))
with profile(activities=[ProfilerActivity.CPU], profile_memory=True) as prof:
filtering.wiener(v, x_as_real)
print("Open-Unmix")
print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=5))
results = []
for ns, ni in product(nb_sources, nb_iterations):
label = 'wiener'
sub_label = f'[{ns}, {ni}]'
v = torch.rand(nb_frames, nb_bins, nb_channels, ns)
x = torch.randn(nb_frames, nb_bins, nb_channels) + 1j * \
torch.randn(nb_frames, nb_bins, nb_channels)
x_as_real = torch.view_as_real(x)
for num_threads in [1, 2, 4]:
results.append(benchmark.Timer(
stmt='y = wiener(v, x, i)',
setup='from norbert.torch import wiener',
globals={'x': x.unsqueeze(0), 'v': v.unsqueeze(0), 'i': ni},
label=label,
num_threads=num_threads,
sub_label=sub_label,
description='norbert torch',
).blocked_autorange(min_run_time=1))
results.append(benchmark.Timer(
stmt='y = wiener(v, x, i)',
setup='from openunmix.filtering import wiener',
globals={'x': x_as_real, 'v': v, 'i': ni},
label=label,
num_threads=num_threads,
sub_label=sub_label,
description='openunmix',
).blocked_autorange(min_run_time=1))
compare = benchmark.Compare(results)
compare.print() |
|
@yoyololicon sorry for the slow response. I will have a look next week 🙏 |
|
This looks great, @faroit would be amazing to be able to pip install this |
|
@yoyololicon Totally forgot about this one. I will have a look again |
|
@faroit No problem. Take your time. Let's first enjoy the ISMIR conference 😆 |
Feature
An alternative and fully differentiable backend of norbert for PyTorch users. One extra dimension is required for input Tensors to make it be able to process in batch, which comes in handy for deep learning training. This implementation had been used to train Danna-Sep 12. This change is backward compatible so it won't affect any existing projects that depend on norbert.
Usage
To use the PyTorch backend, users can simply just replace each call of
norbert.*intonorbert.torch.*.Available Functions
expectation_maximizationwienersoftmaskwiener_gainapply_filterget_mix_modelget_local_gaussian_modelresidual_modelreduce_interferencesNote
In order to make the new backend differentiable, some in-place operations in the original code were re-written so it's not a simple one-to-one translation.
BTW, I have some problems building the docs so I'm not sure what the documentation will be like after this PR. If anyone could help me check the docs I'll appreciate it.
Footnotes
Yu, Chin-Yun, and Kin-Wai Cheuk. "Danna-Sep: Unite to separate them all". https://arxiv.org/abs/2112.03752 ↩
Yuki et al. "Music Demixing Challenge 2021". https://doi.org/10.3389/frsip.2021.808395 ↩