Make the code compatible with other keras backends by scarlehoff · Pull Request #2137 · NNPDF/nnpdf

scarlehoff · 2024-07-26T18:37:06Z

This deals with #2134

It exchanges tf imports to keras (with small changes when the functions signatures change). It was necessary to add a few explicit output shapes here and there so that Keras knows how to compile the graph (and then we can also use the jit_compile flag, at least with the tensorflow backend).

Now, for some comparisons:

In this report master is compared to this branch https://vp.nnpdf.science/ZGmx7ro7SNCESqFEmR2IcA==, TensorFlow backend and 70 replicas running in GPU (A6000) (53 min vs 43 min)
Tensorflow Vs Pytorch (both running in this branch and GPU) https://vp.nnpdf.science/zKyfSvs8TxSI1FXGR72f3w== (43 min vs 178 min)

Turns out that the code is mostly general enough to work with pytorch ootb. Very few changes are needed* (the second commit, and not even, because some of the changes are tf. that I missed).

However, I've run a few quick benchmarks and turns out that it is about 4 times slower than the tensorflow backend. This is not surprising since we have been using tensorflow for a long time so we might be doing thing that just happen to be better for tensorflow. Also, while jit_compile is set to auto, I think the default for pytorch is False.

scarlehoff · 2024-11-22T21:14:37Z

This makes the code:

Keras-only (so no more tensorflow functions in the middle)
Compilable wtih xla (the fit is very much not the bottleneck anymore, but for reference, running in GPU it goes from 20 replicas in 23 minutes to 20 replicas in 18 minutes)
Compatible with pytorch in both cpu and gpu (however running pytorch in graph mode crashed*)
Almost compatible with jax. For jax there are 2 problems: a) the preprocessing layer b) the computation of the losses. a) is fixable but b) is too much work imho so I won't do anything for that unless it turns out to be the best thing ever smh.
And we can now test in conda with python 3.12 :) (fair enough, probably this can be achieved also in master removing the pin for tf 2.17)

*I don't think we want to run with pytorch, but having the possibility is nice. If we start using it for whatever reason I can debug that.

github-actions · 2024-11-22T22:31:52Z

Greetings from your nice fit 🤖 !
I have good news for you, I just finished my tasks:

Fit Name: NNBOT-e4683516b-2024-11-22
Fit Report wrt master: https://vp.nnpdf.science/7V0wqFi7Q4qg4gwkqCTIsg==
Fit Report wrt latest stable reference: https://vp.nnpdf.science/mlZHl8qmQ8eOUmpm9Vc6fg==
Fit Data: https://data.nnpdf.science/fits/NNBOT-e4683516b-2024-11-22.tar.gz

Check the report carefully, and please buy me a ☕ , or better, a GPU 😉!

github-actions · 2024-11-27T17:45:51Z

Greetings from your nice fit 🤖 !
I have good news for you, I just finished my tasks:

Fit Name: NNBOT-50a9d4ec8-2024-11-27
Fit Report wrt master: https://vp.nnpdf.science/uXHKDYSrSG6tjarNT-C4eQ==
Fit Report wrt latest stable reference: https://vp.nnpdf.science/x4BS65iHSXCj7gPH_7AEow==
Fit Data: https://data.nnpdf.science/fits/NNBOT-50a9d4ec8-2024-11-27.tar.gz

Check the report carefully, and please buy me a ☕ , or better, a GPU 😉!

github-actions · 2024-11-29T09:59:23Z

Greetings from your nice fit 🤖 !
I have good news for you, I just finished my tasks:

Fit Name: NNBOT-c75c4f07e-2024-11-29
Fit Report wrt master: https://vp.nnpdf.science/5paW1A3zRhOJNq6wMs2p2w==
Fit Report wrt latest stable reference: https://vp.nnpdf.science/v0MST-dHR66LAfUsGvileg==
Fit Data: https://data.nnpdf.science/fits/NNBOT-c75c4f07e-2024-11-29.tar.gz

Check the report carefully, and please buy me a ☕ , or better, a GPU 😉!

scarlehoff · 2024-11-29T13:45:45Z

Ok, this is ready for review. Some comparisons in the first post.

I've heard you are completely free this weekend @RoyStegeman in case you want to have another look :P

RoyStegeman

Thanks, looks good. Main point is that it should be documented

make the code compatible with pytorch make it more keras only, working also with sum rules update keras limits test with newer tf fix for normalized distributions

test change workflow test in 3.12 in conda

remove other instances of kops Update n3fit/src/n3fit/backends/keras_backend/MetaModel.py Co-authored-by: Roy Stegeman <roystegeman@live.nl>

…ckend

scarlehoff added the n3fit Issues and PRs related to n3fit label Jul 26, 2024

scarlehoff marked this pull request as draft July 26, 2024 18:47

RoyStegeman force-pushed the disable_xla_tf216 branch from 5721909 to 2120355 Compare July 27, 2024 14:39

Base automatically changed from disable_xla_tf216 to master July 28, 2024 11:12

scarlehoff mentioned this pull request Jul 29, 2024

Disable jit compilation in tf > 2.16 #2135

Merged

scarlehoff force-pushed the make_code_tf_independent branch from f2484d3 to 4448ba4 Compare August 16, 2024 09:38

scarlehoff force-pushed the make_code_tf_independent branch 3 times, most recently from 502ac27 to 2f150e6 Compare November 22, 2024 14:42

scarlehoff changed the title ~~[WIP] Make the code compatible with other keras backends~~ Make the code compatible with other keras backends Nov 22, 2024

scarlehoff marked this pull request as ready for review November 22, 2024 16:20

scarlehoff added the run-fit-bot Starts fit bot from a PR. label Nov 22, 2024

scarlehoff force-pushed the make_code_tf_independent branch 2 times, most recently from 45fe562 to b3c5b27 Compare November 22, 2024 21:11

scarlehoff added run-fit-bot Starts fit bot from a PR. and removed run-fit-bot Starts fit bot from a PR. labels Nov 22, 2024

scarlehoff requested a review from RoyStegeman November 22, 2024 22:54

RoyStegeman force-pushed the make_code_tf_independent branch from 6b94412 to b859c07 Compare November 25, 2024 14:47

RoyStegeman reviewed Nov 25, 2024

View reviewed changes

scarlehoff force-pushed the make_code_tf_independent branch 2 times, most recently from 59464e9 to ee35538 Compare November 27, 2024 16:34

scarlehoff changed the base branch from master to rerun_fitbot_241127 November 27, 2024 16:35

scarlehoff added run-fit-bot Starts fit bot from a PR. and removed run-fit-bot Starts fit bot from a PR. labels Nov 27, 2024

Base automatically changed from rerun_fitbot_241127 to master November 27, 2024 18:20

scarlehoff force-pushed the make_code_tf_independent branch from ee35538 to bd05e0f Compare November 27, 2024 18:20

scarlehoff mentioned this pull request Nov 28, 2024

Run fitbot with pip and not conda #2232

Merged

scarlehoff removed the run-fit-bot Starts fit bot from a PR. label Nov 28, 2024

scarlehoff force-pushed the make_code_tf_independent branch from 3ca453d to 15e9f34 Compare November 29, 2024 07:28

scarlehoff added the run-fit-bot Starts fit bot from a PR. label Nov 29, 2024

scarlehoff requested review from Radonirinaunimi and RoyStegeman November 29, 2024 13:45

scarlehoff force-pushed the make_code_tf_independent branch from 15e9f34 to 23e7620 Compare December 4, 2024 15:36

RoyStegeman force-pushed the make_code_tf_independent branch from 23e7620 to da1d6bd Compare December 5, 2024 10:21

RoyStegeman linked an issue Dec 5, 2024 that may be closed by this pull request

Move from tensorflow.keras to keras #2134

Closed

RoyStegeman reviewed Dec 5, 2024

View reviewed changes

RoyStegeman force-pushed the make_code_tf_independent branch from f97a928 to af832de Compare December 5, 2024 13:59

scarlehoff and others added 7 commits December 6, 2024 12:22

remove most tf dependencies

92d349d

make the code compatible with pytorch make it more keras only, working also with sum rules update keras limits test with newer tf fix for normalized distributions

add a test for pytorch

fd85990

test change workflow test in 3.12 in conda

use xla to compile the tf graph, add _some_ jax support

82e4b73

apply some comments, refactor operations

909424f

remove other instances of kops Update n3fit/src/n3fit/backends/keras_backend/MetaModel.py Co-authored-by: Roy Stegeman <roystegeman@live.nl>

apply comments ; add set_threads from pytorch as well

9e7ae84

more review comments

eaec1c8

update docs ; update versions at the end of the fit to include the ba…

59ce25a

…ckend

RoyStegeman force-pushed the make_code_tf_independent branch from 300c865 to 59ce25a Compare December 6, 2024 12:22

RoyStegeman approved these changes Dec 6, 2024

View reviewed changes

Comment thread doc/sphinx/source/n3fit/methodology.rst Outdated

Comment thread doc/sphinx/source/tutorials/run-fit.rst Outdated

RoyStegeman added 2 commits December 6, 2024 12:28

Update doc/sphinx/source/n3fit/methodology.rst

4cfba2c

Update doc/sphinx/source/tutorials/run-fit.rst

bd17587

RoyStegeman merged commit f4467d0 into master Dec 6, 2024

RoyStegeman deleted the make_code_tf_independent branch December 6, 2024 13:29

scarlehoff mentioned this pull request Dec 6, 2024

To Do for 4.0.10 #1854

Closed

33 tasks

Conversation

scarlehoff commented Jul 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

scarlehoff commented Nov 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Nov 22, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Nov 27, 2024

Uh oh!

github-actions Bot commented Nov 29, 2024

Uh oh!

scarlehoff commented Nov 29, 2024

Uh oh!

RoyStegeman left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

scarlehoff commented Jul 26, 2024 •

edited

Loading

scarlehoff commented Nov 22, 2024 •

edited

Loading