diff --git a/.github/workflows/pytorch_test.yml b/.github/workflows/pytorch_test.yml new file mode 100644 index 0000000000..ca20ca6fa2 --- /dev/null +++ b/.github/workflows/pytorch_test.yml @@ -0,0 +1,28 @@ +name: Test the python installation + +on: [push] + +jobs: + build: + strategy: + matrix: + os: [ubuntu-latest] + python-version: ["3.12"] + include: + - os: ubuntu-latest + fail-fast: false + runs-on: ${{ matrix.os }} + env: + NETRC_FILE: ${{ secrets.NETRC_FILE }} + NNPDF_SSH_KEY: ${{ secrets.NNPDF_SSH_KEY }} + steps: + - uses: actions/checkout@v4 + - name: Install nnpdf without LHAPDF + shell: bash -l {0} + run: | + pip install .[nolha,torch] + - name: Test we can run one runcard + shell: bash -l {0} + run: | + cd n3fit/runcards/examples + n3fit Basic_runcard.yml 4 diff --git a/n3fit/src/n3fit/backends/keras_backend/MetaModel.py b/n3fit/src/n3fit/backends/keras_backend/MetaModel.py index f1cdbc418a..afcc4b6dad 100644 --- a/n3fit/src/n3fit/backends/keras_backend/MetaModel.py +++ b/n3fit/src/n3fit/backends/keras_backend/MetaModel.py @@ -8,12 +8,11 @@ from pathlib import Path import re -from keras import backend as K +from keras import Variable from keras import ops as Kops from keras import optimizers as Kopt from keras.models import Model import numpy as np -import tensorflow as tf import n3fit.backends.keras_backend.operations as op @@ -501,9 +500,9 @@ def get_layer_replica_weights(layer, i_replica: int): """ if is_stacked_single_replicas(layer): weights_ref = layer.get_layer(f"{NN_PREFIX}_{i_replica}").weights - weights = [tf.Variable(w, name=w.name) for w in weights_ref] + weights = [Variable(w, name=w.name) for w in weights_ref] else: - weights = [tf.Variable(w[i_replica : i_replica + 1], name=w.name) for w in layer.weights] + weights = [Variable(w[i_replica : i_replica + 1], name=w.name) for w in layer.weights] return weights diff --git a/n3fit/src/n3fit/backends/keras_backend/callbacks.py b/n3fit/src/n3fit/backends/keras_backend/callbacks.py index 29324e7cee..8a593af814 100644 --- a/n3fit/src/n3fit/backends/keras_backend/callbacks.py +++ b/n3fit/src/n3fit/backends/keras_backend/callbacks.py @@ -15,9 +15,11 @@ import logging from time import time +# Callbacks need tensorflow installed even if the backend is pytorch from keras.callbacks import Callback, TensorBoard import numpy as np -import tensorflow as tf + +from .operations import decorator_compiler log = logging.getLogger(__name__) @@ -171,6 +173,7 @@ def on_train_begin(self, logs=None): layer = self.model.get_layer(layer_name) self.updateable_weights.append(layer.weights) + @decorator_compiler def _update_weights(self): """Update all the weight with the corresponding multipliers Wrapped with tf.function to compensate the for loops as both weights variables diff --git a/n3fit/src/n3fit/backends/keras_backend/internal_state.py b/n3fit/src/n3fit/backends/keras_backend/internal_state.py index dd61068190..4235a73473 100644 --- a/n3fit/src/n3fit/backends/keras_backend/internal_state.py +++ b/n3fit/src/n3fit/backends/keras_backend/internal_state.py @@ -17,17 +17,34 @@ import keras from keras import backend as K import numpy as np -import tensorflow as tf log = logging.getLogger(__name__) +# Prepare Keras-backend dependent functions +if K.backend() == "torch": -def set_eager(flag=True): - """Set eager mode on or off - for a very slow but fine grained debugging call this function as early as possible - ideally after the first tf import - """ - tf.config.run_functions_eagerly(flag) + def set_eager(flag=True): + """Pytorch is eager by default""" + pass + + def set_threading(threads, core): + """Not implemented""" + pass + +elif K.backend() == "tensorflow": + import tensorflow as tf + + def set_eager(flag=True): + """Set eager mode on or off + for a very slow but fine grained debugging call this function as early as possible + ideally after the first tf import + """ + tf.config.run_functions_eagerly(flag) + + def set_threading(threads, cores): + """Set the Tensorflow inter and intra parallelism options""" + tf.config.threading.set_inter_op_parallelism_threads(threads) + tf.config.threading.set_intra_op_parallelism_threads(cores) def set_number_of_cores(max_cores=None, max_threads=None): @@ -65,8 +82,7 @@ def set_number_of_cores(max_cores=None, max_threads=None): log.info("Setting the number of cores to: %d", cores) try: - tf.config.threading.set_inter_op_parallelism_threads(threads) - tf.config.threading.set_intra_op_parallelism_threads(cores) + set_threading(threads, cores) except RuntimeError: # If pdfflow is being used, tensorflow will already be initialized by pdfflow # maybe it would be good to drop completely pdfflow before starting the fit? (TODO ?) @@ -130,7 +146,7 @@ def set_initial_state(debug=False, external_seed=None, max_cores=None, double_pr clear_backend_state() if double_precision: - tf.keras.backend.set_floatx('float64') + K.set_floatx('float64') # Set the number of cores depending on the user choice of max_cores # if debug mode and no number of cores set by the user, set to 1 @@ -143,7 +159,11 @@ def set_initial_state(debug=False, external_seed=None, max_cores=None, double_pr # Once again, if in debug mode or external_seed set, set also the TF seed if debug or external_seed: - tf.random.set_seed(use_seed) + if K.backend() == "tensorflow": + # if backend is tensorflow, keep the old seeding + tf.random.set_seed(use_seed) + else: + keras.utils.set_random_seed(use_seed) def get_physical_gpus(): diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py index f521b0536e..ae65e1f1de 100644 --- a/n3fit/src/n3fit/backends/keras_backend/operations.py +++ b/n3fit/src/n3fit/backends/keras_backend/operations.py @@ -122,7 +122,7 @@ def batchit(x, batch_dimension=0, **kwarg): # layer generation -def numpy_to_input(numpy_array: np.typing.NDArray, name: Optional[str] = None): +def numpy_to_input(numpy_array, name=None): """ Takes a numpy array and generates an Input layer with the same shape, but with a batch dimension (of size 1) added. diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py index a66e03a3fe..721eb6b38d 100644 --- a/n3fit/src/n3fit/msr.py +++ b/n3fit/src/n3fit/msr.py @@ -84,7 +84,9 @@ def generate_msr_model_and_grid( # 3. Prepare the pdf for integration by dividing by x pdf_integrand = Lambda( - lambda x_pdf: op.batchit(x_pdf[0], batch_dimension=1) * x_pdf[1], name="pdf_integrand" + lambda x_pdf: op.batchit(x_pdf[0], batch_dimension=1) * x_pdf[1], + name="pdf_integrand", + output_shape=pdf_xgrid_integration.shape[1:], )([x_divided, pdf_xgrid_integration]) # 4. Integrate the pdf diff --git a/pyproject.toml b/pyproject.toml index 28b19d67cb..8058a437f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,8 @@ fiatlux = {version = "*", optional = true} # without lhapdf pdfflow = {version = "^1.2.1", optional = true} lhapdf-management = {version = "^0.5", optional = true} +# torch +torch = {version = "*", optional = true} # Optional dependencies [tool.poetry.extras] @@ -104,6 +106,7 @@ tests = ["pytest", "pytest-mpl", "hypothesis"] docs = ["sphinxcontrib", "sphinx-rtd-theme", "sphinx", "tabulate"] qed = ["fiatlux"] nolha = ["pdfflow", "lhapdf-management"] +torch = ["torch"] [tool.poetry-dynamic-versioning] enable = true