Multi-stem tensor fix, input length=output length option, MLP depth adjustment

[crlandsc]

Changes:

• Added option to output a tensor the same length as input tensor.
  • Implemented via padding in the torch.istft function.
  • When utilizing an external method for calculating loss, this helps when comparing input & output audio.
• Tensor shape fixes for multi-stem training.
  • Training on a single stem worked, but I ran into tensor shape issues when I tried to add stems.
• Removed subtraction in MLP dims calculation to align mask_estimator_depth and MLP depth to be equal to the number of MLP hidden layers.
  • The mask_estimator_depth=2 attribute seemed misleading when it resulted in an MLP depth of 3 layers (1 in, 1 hidden, 1 out). The default value for the MLP depth was 1 and mask_estimator_depth was 2. Now these are aligned to both be 1 and the attribute directly corresponds to the number of hidden layers in the MLPs.
  • Also, I believe that the MLPs only have 1 hidden layer per the BS-RoFormer paper (please correct me if I am incorrect!).

[lucidrains]

@crlandsc hey Christopher! this looks great! thank you for addressing the multiple stem situation for mel-band roformer; overlooked it

[crlandsc]

No problem, thanks for reviewing so quickly and for this paper implementation! It's fantastic!

Are all of the default args for the mel-roformer based on the paper? Because I was having trouble replicating the exact number of parameters that they detail.

[lucidrains]

thanks for the kind words!

there could be tiny differences, but i think all the main points of the paper are already in there (axial attention, rotary relative positions, and overlapping frequencies bands for mel-band roformer variant)

[lucidrains]

by default i am using the librosa mel bank as they detailed in the paper

[crlandsc]

Ok, that is pretty much what I thought. I was also making sure that I wasn't messing something up on my end as I started to train with it. Thanks again!

[lucidrains]

thanks for the PR!