Skip to content

Commit

Permalink
scale test on lumi, fix horovod (#189)
Browse files Browse the repository at this point in the history 
* scale test on lumi

* fix horovod
  • Loading branch information
@jpata
jpata authored Aug 10, 2023
1 parent ae5ec59 commit 17875fe
Show file tree
Hide file tree
Showing 25 changed files with 847 additions and 127 deletions.
62 changes: 62 additions & 0 deletions mlpf/heptfds/clic_pf_edm4hep_hits/single_ele.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
from pathlib import Path

import tensorflow as tf
from utils_edm import (
X_FEATURES_CH,
X_FEATURES_TRK,
Y_FEATURES,
generate_examples,
split_sample_several,
)

import tensorflow_datasets as tfds

_DESCRIPTION = """
CLIC EDM4HEP dataset with single electron with raw hits
"""

_CITATION = """
"""


class ClicEdmSingleElectronHitsPf(tfds.core.GeneratorBasedBuilder):
VERSION = tfds.core.Version("1.2.0")
RELEASE_NOTES = {"1.1.0": "Remove track referencepoint feature", "1.2.0": "Keep all interacting genparticels"}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
FIXME
"""

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(
shape=(
None,
max(len(X_FEATURES_TRK), len(X_FEATURES_CH)),
),
dtype=tf.float32,
),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=None,
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(
x_features_track=X_FEATURES_TRK,
x_features_calohit=X_FEATURES_CH,
y_features=Y_FEATURES,
),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
path = dl_manager.manual_dir
return split_sample_several([Path(path / "e-/"), Path(path / "e+/")])

def _generate_examples(self, files):
return generate_examples(files)
65 changes: 65 additions & 0 deletions mlpf/heptfds/clic_pf_edm4hep_hits/single_gamma.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
from pathlib import Path

import tensorflow as tf
from utils_edm import (
X_FEATURES_CH,
X_FEATURES_TRK,
Y_FEATURES,
generate_examples,
split_sample,
)

import tensorflow_datasets as tfds

_DESCRIPTION = """
CLIC EDM4HEP dataset with single gamma with raw hits
"""

_CITATION = """
"""


class ClicEdmSingleGammaHitsPf(tfds.core.GeneratorBasedBuilder):
VERSION = tfds.core.Version("1.2.0")
RELEASE_NOTES = {
"1.1.0": "Remove track referencepoint feature",
"1.2.0": "Keep all interacting genparticles",
}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
FIXME
"""

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(
shape=(
None,
max(len(X_FEATURES_TRK), len(X_FEATURES_CH)),
),
dtype=tf.float32,
),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=None,
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(
x_features_track=X_FEATURES_TRK,
x_features_calohit=X_FEATURES_CH,
y_features=Y_FEATURES,
),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
path = dl_manager.manual_dir
return split_sample(Path(path / "gamma/"))

def _generate_examples(self, files):
return generate_examples(files)
65 changes: 65 additions & 0 deletions mlpf/heptfds/clic_pf_edm4hep_hits/single_kaon0L.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
from pathlib import Path

import tensorflow as tf
from utils_edm import (
X_FEATURES_CH,
X_FEATURES_TRK,
Y_FEATURES,
generate_examples,
split_sample,
)

import tensorflow_datasets as tfds

_DESCRIPTION = """
CLIC EDM4HEP dataset with single kaon0L with raw hits
"""

_CITATION = """
"""


class ClicEdmSingleKaon0lHitsPf(tfds.core.GeneratorBasedBuilder):
VERSION = tfds.core.Version("1.2.0")
RELEASE_NOTES = {
"1.1.0": "Remove track referencepoint feature",
"1.2.0": "Keep all interacting genparticles",
}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
FIXME
"""

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(
shape=(
None,
max(len(X_FEATURES_TRK), len(X_FEATURES_CH)),
),
dtype=tf.float32,
),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=None,
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(
x_features_track=X_FEATURES_TRK,
x_features_calohit=X_FEATURES_CH,
y_features=Y_FEATURES,
),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
path = dl_manager.manual_dir
return split_sample(Path(path / "kaon0L/"))

def _generate_examples(self, files):
return generate_examples(files)
62 changes: 62 additions & 0 deletions mlpf/heptfds/clic_pf_edm4hep_hits/single_mu.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
from pathlib import Path

import tensorflow as tf
from utils_edm import (
X_FEATURES_CH,
X_FEATURES_TRK,
Y_FEATURES,
generate_examples,
split_sample_several,
)

import tensorflow_datasets as tfds

_DESCRIPTION = """
CLIC EDM4HEP dataset with single muon with raw hits
"""

_CITATION = """
"""


class ClicEdmSingleMuonHitsPf(tfds.core.GeneratorBasedBuilder):
VERSION = tfds.core.Version("1.2.0")
RELEASE_NOTES = {"1.1.0": "Remove track referencepoint feature", "1.2.0": "Keep all interacting genparticels"}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
FIXME
"""

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(
shape=(
None,
max(len(X_FEATURES_TRK), len(X_FEATURES_CH)),
),
dtype=tf.float32,
),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=None,
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(
x_features_track=X_FEATURES_TRK,
x_features_calohit=X_FEATURES_CH,
y_features=Y_FEATURES,
),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
path = dl_manager.manual_dir
return split_sample_several([Path(path / "mu-/"), Path(path / "mu+/")])

def _generate_examples(self, files):
return generate_examples(files)
65 changes: 65 additions & 0 deletions mlpf/heptfds/clic_pf_edm4hep_hits/single_neutron.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
from pathlib import Path

import tensorflow as tf
from utils_edm import (
X_FEATURES_CH,
X_FEATURES_TRK,
Y_FEATURES,
generate_examples,
split_sample,
)

import tensorflow_datasets as tfds

_DESCRIPTION = """
CLIC EDM4HEP dataset with single neutron with raw hits
"""

_CITATION = """
"""


class ClicEdmSingleNeutronHitsPf(tfds.core.GeneratorBasedBuilder):
VERSION = tfds.core.Version("1.2.0")
RELEASE_NOTES = {
"1.1.0": "Remove track referencepoint feature",
"1.2.0": "Keep all interacting genparticles",
}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
FIXME
"""

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(
shape=(
None,
max(len(X_FEATURES_TRK), len(X_FEATURES_CH)),
),
dtype=tf.float32,
),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=None,
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(
x_features_track=X_FEATURES_TRK,
x_features_calohit=X_FEATURES_CH,
y_features=Y_FEATURES,
),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
path = dl_manager.manual_dir
return split_sample(Path(path / "neutron/"))

def _generate_examples(self, files):
return generate_examples(files)
65 changes: 65 additions & 0 deletions mlpf/heptfds/clic_pf_edm4hep_hits/single_pi.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
from pathlib import Path

import tensorflow as tf
from utils_edm import (
X_FEATURES_CH,
X_FEATURES_TRK,
Y_FEATURES,
generate_examples,
split_sample_several,
)

import tensorflow_datasets as tfds

_DESCRIPTION = """
CLIC EDM4HEP dataset with single pi- with raw hits
"""

_CITATION = """
"""


class ClicEdmSinglePiHitsPf(tfds.core.GeneratorBasedBuilder):
VERSION = tfds.core.Version("1.2.0")
RELEASE_NOTES = {
"1.1.0": "Remove track referencepoint feature",
"1.2.0": "Keep all interacting genparticles",
}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
FIXME
"""

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(
shape=(
None,
max(len(X_FEATURES_TRK), len(X_FEATURES_CH)),
),
dtype=tf.float32,
),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=None,
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(
x_features_track=X_FEATURES_TRK,
x_features_calohit=X_FEATURES_CH,
y_features=Y_FEATURES,
),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
path = dl_manager.manual_dir
return split_sample_several([Path(path / "pi-/"), Path(path / "pi+/")])

def _generate_examples(self, files):
return generate_examples(files)
Loading

0 comments on commit 17875fe

Please sign in to comment.