PyTorch feature parity

[CarloLucibello]

A list of PyTorch 1.7 features.
Items are checked if we have something more or less equivalent in Flux or in the julia ecosystem and supported by Flux.
This list is not complete, it comes from a rough scan of pytorch's documentation. Please feel free to add anything I missed in the comments, and whoever has write access to modify the list.
Related issue https://github.com/FluxML/ML-Coordination-Tracker/issues/16, and more generally anything in https://github.com/FluxML/ML-Coordination-Tracker/issues

Pytorch Features

Conv Layers

• Conv1d, Conv2d, Conv3d.
• ConvTranspose1d, ConvTranspose2d, ConvTranspose3d.
• groups in convolution layers
• Fold, Unfold. In progress: Add fold and unfold NNlib.jl#444

Pooling Layers

• MaxPool1d, MaxPool2d, MaxPool3d
• MaxUnPool1d, MaxUnPool2d, MaxUnPool3d
• AvgPool1d, AvgPool2d, AvgPool3d
• FractionalMaxPool2d
• LPPool1d, LPPool2d
• AdaptiveAvgPool1d, AdaptiveAvgPool2d, AdaptiveAvgPool3d
• AdaptiveMaxPool1d, AdaptiveMaxPool2d, AdaptiveMaxPool3d

Padding Layers

• ReflectionPad (1d,2d)
• ReplicationPad (1d,2d,3d) ( NNlib.pad_repeat)
• ZeroPad (2d)
• ConstantPad (1d,2d,3d)
• Add corresponding layers for all of the aboves wrapping the NNlin functions keep as functions. Need to add them Flux's docs.

Activations

• ... . NNlib has an extensive collection of activation, plus we have any julia function.

Normalization Layers

• BatchNorm1d, BatchNorm2d, BatchNorm3d
• LayerNorm
• GroupNorm
• InstanceNorm1d,InstanceNorm2d,InstanceNorm3d
• SyncBatchNorm
• LocalResponseNorm. Very old unfinished PR Local Response Normalization[W.I.P] #312. It is an outdated technique, probably we can live without it.
• Move the functional implementations to NNlib.jl (BatchNorm and Dropout NNlib.jl#19)

Recurrent Layers

• RNN
• GRU
• LSTM

Attention Layers

• Transformer. Well maintained implementations in Tansformers.jl.
• MultiHeadAttention Should be moved from Transformers.jl to Flux.jl (ensure hitting cudnn kernels). PR MultiHeadAttention implementation #2146

Linear Layers

• Identity
• Linear
• Bilinear

Dropout Layers

• Dropout
• Dropout2d, Dropout3d (Make Dropout docstring clear w.r.t. N-D dropout #1490)
• AlphaDropout

Sparse Layers

• Embedding PR add Embedding layer #1516
• EmbeddingBag PR Add EmbeddingBag #2031

Distance Functions

• CosineSimilarity. We have this in Distances.jl. Also easy to handcode. TODO check if AD and gpu friendly.
• PairwiseDistance. We have this in Distances.jl TODO check if AD and gpu friendly (could use Tullio.jl to achieve both)

Loss Functions

• .... . We should be well covered here.
• CTCLoss. Being Implemented in Add CTC loss to new Losses module #1287 (todo: remove separate GPU case, integrate with cudnn)

Vision Layers

• PixelShuffle. add Upsample and PixelShuffle layers #1468
• Upsample (for 1d, 2d, and 3d). (partially done in add Upsample and PixelShuffle layers #1468)
  • 'nearest'
  • 'linear' (cpu version merged in NNlib, CUDA PR still to come)
  • 'bilinear'
  • 'bicubic'
  • 'trilinear' (cpu versino merged in NNlib, CUDA PR still open )

Initialization

• xavier_uniform, xavier_normal. Called glorot here.
• kaiming_normal kaiming_uniform
• sparse
• orthogonal (Add Orthogonal initialization feature. #1496)

Parallelism and Distributed

• DataParallel
• DistributedDataParallel(solved by https://github.com/DhairyaLGandhi/DaggerFlux.jl
• set_num_threads, set_num_interop_threads. Not sure which operations are parallelized in pytorch. Here we have parallelization only in blas operations.

Distributions

• diff rules for logpdf offered by DistributionsAD.jl
• rsample. params's differentiability through sampling supported by many distr: gradient(mu -> rand(Normal(mu, 1)), 0) == (1,).

ONNX

• Current best support in ONNXmutable. See this discussion
  • ONNX.jl's old implementation has been replaced
  • Overcome the limitations reported here

FFT

• ... . Zygote has the adjoints for AbstractFFTs.

Quantization

• ...

Pruning

• WIP pruning package here

Optim

• schedulers Add simple Schedulers #1434 and Add basic scheduling policies and a scheduler #1506, also see ParameterSchedulers.jl
  • Integrate with Flux's optimizers? (See RFC: basic sketch of scheduling Optimisers.jl#15)
  • Document in Flux (see Add ParameterSchedulers.jl to docs #1511 and Update for latest ParameterSchedulers.jl release #1513)
    - [ ] Reexport in Flux (see Add basic scheduling policies and a scheduler #1506) (TBD)
  • LambdaLR (handled in ParameterSchedulers.jl)
  • MultiplicativeLR (handled in ParameterSchedulers.jl)
• optimizers
  • SGD (+ momentum)
  • Adam
  • AdaGrad
  • AdaDelta
  • RMSprop
  • LBFGS. Integration with Optim.jl

LinAlg

• det
• norm

Tensorboard

• integration offered by TensorBoardLogger.jl

XLA

• Some work in XLA.jl

Misc

• Pytorch has both layers and their functional counterpart.
• einsum. AD and CUDA compatible Einstein summation given by Tullio.jl and other packages
  • add Documentation to Flux.jl
• LazyModuleMixin (pytorch 1.8) PR Add @autosize #2078
• weight_norm. Attempt in Added WeightNorm #1005 , PR Add WeightNorm layer #2053
• modules iterator. define modules function #1444
• spectral_norm. Old attempt in Fixed the spectral normalization #115

Pytorch Extras

Torchvision

• datasets. Some are implemented in DLDatasets.jl (unreleased), some in FastAI.jl, some in MLDatasets.jl, many are missing.
  • Will consolidate in MLDatasets.jl (see merge into MLDatasets? lorenzoh/DLDatasets.jl#1)
• models. Some are implemented in Metalhead.jl, but it is a bit stale and not comprehensive.
  • Metalhead's PR should add a bunch of model and generally revive the repo
  • We should expose the possibility to load pretrained weights
• io
• transforms. Some unreleased work in DataAugmentation.jl

Torchaudio
...

Torchtext
...

[gxyd]

Do you mind if I try to implement the support in Flux corresponding to Dropout2D in pytorch?

[CarloLucibello]

yes please, essentially this is all up for grabs

[DhairyaLGandhi]

Note that we shouldn't add all these layers here, for eg pixel shuffle has an implementation, so does Transformers, up sampling and embedding are direct Julia operations etc

[CarloLucibello]

• where is pixel shuffle?
• upsampling is not trivial, we already have a few unfinished attempts here (Add bilinear upsample layer #1180 )
• I was surprised to find transformers in pytorch, they were requested here [feature request] nn.Transformer pytorch/pytorch#10459, I guess it makes sense though to have some basic components in Flux in the same way we have RNN.
• I'm not an NLP guy, but it looks like learnable embeddings https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html are not entirely trivial to implement in a way that plays nicely with other layers. I think they are very very common, so should be implemented in Flux.

Maybe @chengchingwen could provide some suggestions on the last two items

[DrChainsaw]

@CarloLucibello About ONNX: This exists which I guess (hope?) is better than nothing: https://github.com/DrChainsaw/ONNXmutable.jl

I haven't registered it because 1) the name sucks and I can't think of anything better and 2) I'm thinking of splitting the import and export into two separate packages. 1 is the main blocker though :)

I'd be happy to donate it to FluxML, or parts of it (e.g. import/export primitives).

[darsnack]

Yeah upsampling is non-trivial to get right and be performant on the GPU as well (last time I tried it, I had to ask in #gpu on Slack to get a good implementation).

For ONNX, is it possible to hand control of ONNX.jl to @DrChainsaw? It seems like ONNXmutable.jl should really supersede that package.

For vision models, there is this Metalhead PR which I think will bring us much closer to PyTorch parity. I am planning on training some of the simpler ones this weekend, but I would appreciate the help to add pre-trained weights from anyone with a GPU.

Lastly, for hyperparameter/learning rate schedules, I just started ParameterSchedulers.jl to break the functionality out of FluxTraining.jl. This is quite a simple package, and I want to finish it this weekend for a project. I am happy to transfer ownership to FluxML.

[bhvieira]

I tried implementing WeightNorm before, but it's harder than I thought without doing a per-layer implementation. See #1005
Doing a per layer implementation is actually easy, but maintenance hell at the same time.

[CarloLucibello]

@DrChainsaw what are the limitations of ONNXmutable?

[DrChainsaw]

@CarloLucibello From the ML-Coordination issue it seems like there is alot of ways to look at ONNX import/export so what is a limitation appears to be a bit more subjective than I thought.

Here are some things I can think of

Only a subset of OPs supported. This is imo not a big deal as I have made an effort for it to be easy to add more and even make it easy for users to just hack in own versions locally. Most OPs are trivial to add but I have intentionally not added more than what I happen to need in hopes that it would encourage contribution.

It has capabilities which perhaps only a small subset of users have use for w.r.t model manipulation. This translates to dependencies like JuMP and Cbc (used to solve the problem of keeping all parameter shapes aligned when changing the model structure) as well as metadata used to formulate the shape constraints. This may appear as bloat to users who only want to import a model and use it. The annoying part here is that Chain can't represent any graph and even things like what is proposed in #1289 seem very hard to translate to from a more standard graph format such as the one used in ONNX. NaiveNASlib has an internal graph format which does not have the extra functionality for shape alignment which could perhaps be used, but there seem to be a desire for a 'Flux native' format.

RNNs are currently a bit limited, although this is more on Flux than on ONNXmutable since ONNX wants RNN to have 3D input while Flux wants 2D (in a loop). I have worked around this to some extent by just changing the model shape to 3D if a recurrent layer is found and then just fold the time dimension into the batch dimension if a Dense layer is encountered. This only works for a few model architecture types though.

Exporting functionality can't handle 1) non-primitive functions with type constraints (e.g. function thewholemodel(x::AbstractArray)) and 2) non-function control flow (e.g if/else/for, functions like ifelse/map/reduce or let_onnx_know_this_is_a_loop(f, n) could be solveable I think). The first can probably be hacked around with IRtools but I think that the latter would require some abstract interpreter or similar sopisticated code analysis/transformation, e.g. mjolnir.

Eco-system wise it would be better to refactor at least the export primitives to use NNlib as that would make them useable from other libraries which use NNlib (KNet, Avalorn etc). Perhaps not so much a limitation in itself though and can always be broken out later down the road. For export there is no limit on how many ways one can chose to translate a Julia function to an ONNX Node.

Btw, I think it would be better to try to remove ONNX.jl from the general registry and use a name like OnnxFlux.jl to clearly state that it translates between ONNX and Flux.

[darsnack]

Btw, I think it would be better to try to remove ONNX.jl from the general registry and use a name like OnnxFlux.jl to clearly state that it translates between ONNX and Flux.

Unfortunately we can't remove packages from the registry. But if ONNXFlux.jl makes more sense, then we can just archive the ONNX.jl repo.

[ToucheSir]

I don't think it's unreasonable to expect anyone looking to use transformer layers to use Transformers.jl. One potential reason for Torch to add them is because there is no canonical library for transformers in that ecosystem (or really for any other domain...).

RE ONNX, why not give that repo name over to ONNXMutable and then consider how best to refactor/reorganize? I highly doubt anyone is using the existing functionality, given that it's broken on most recent versions of Julia that Flux supports.

RE XLA, I presume this is covered by the work Keno and Tim are doing? Not sure if there's a link to any details there.

[jeremiedb]

Regarding embeddings, although I haven't dealt with the potential caveats from weight norm and such, are there challenges I'm overlooking compared to doing a fairly trivial matrix indexing? Example:

struct Embed{T}
    w::T
end

@functor Embed
Embed(in::Integer, out::Integer; initW=glorot_uniform) = Embed(initW(out, in))
(m::Embed)(x::AbstractVector) = m.w[:,x]

[ToucheSir]

My understanding is that the trivial indexing triggers scalar indexing on GPU arrays. Transformers.jl has custom implementations for both CPU and CUDA, so in that sense the hard work is already done.

[ToucheSir]

Something else I'd like to submit for consideration is an equivalent to the upcoming LazyModuleMixin. Not a 1-1 port, but some mechanism to avoid specifying intermediate sizes during model construction.