Jammy flow integration #728
Conversation
|
@Aske-Rosted checks are now passing. Please let me know if you have any questions |
| sensor listed here will be removed from the graph. Defaults to None. | ||
| string_mask: A list of string id's to be masked from the graph. Defaults to None. | ||
| sensor listed here will be removed from the graph. | ||
| Defaults to None. |
There was a problem hiding this comment.
Did some unintended formatting happen here?
| return self._default_prediction_labels | ||
|
|
||
| def nb_inputs(self) -> Union[int, None]: # type: ignore | ||
| """Return number of conditional inputs assumed by task.""" |
There was a problem hiding this comment.
does non-conditional inputs also exist or what is the reason for the distinction of "conditional"?
There was a problem hiding this comment.
The flow can be trained both as a conditional flow, where it's conditioned on a latent representation of an event e.g. p(target|latent_rep) which is the probability of seeing target for event represented by latent_rep, and as an unconditional probability distribution e.g. p(target) which would just learn the overall distribution of the variable in the dataset.
src/graphnet/models/easy_model.py
Outdated
There was a problem hiding this comment.
Are these changes unrelated to jammy flows?
There was a problem hiding this comment.
reread the minor changes and these changes are described there, you can disregard this question.
src/graphnet/models/graphs/graphs.py
Outdated
There was a problem hiding this comment.
Are these changes also unrelated to Jammy flows?
There was a problem hiding this comment.
reread the minor changes and these changes are described there, you can disregard this question.
src/graphnet/models/task/task.py
Outdated
| labels = labels.to(self.dtype) | ||
| # Set the initial parameters of flow close to truth | ||
| # This speeds up training and helps with NaN | ||
| if self._initialized is False: |
There was a problem hiding this comment.
Just a quick note: I noticed that this part could override the learned weights whenever they are loaded from a state dict, which might affect further training or inference. It could be worth adding a global boolean to check if the weights have been loaded. Otherwise, it might be good at least to pass and self.training , which should work when the model is set to evaluation mode for predictions.
There was a problem hiding this comment.
@jvaracarbonell thanks for pointing this out! I've added the check
…oe/graphnet into jammy_flow_integration
This PR adds support for Normalizing Flows via the jammy_flows package, and therefore supersedes #649. The benefit of using
jammy_flowsis that it contains many different normalizing flows, and that we avoid maintaining that code ourselves :-).The package is not listed as a direct dependency but used as an optional support package.
Specifically, this PR introduces the following major changes:
StandardFlowTasknow uses jammy_flows to construct pdfs of any kind that it supports. These pdfs can be both conditional and non-conditional. Conditional flows can be conditioned on latent model output, event-level information or pulse-level information.NormalizingFlowwhich work with theStandardFlowTask. Usage is similar toStandardModel.examples/04_training/07_train_normalizing_flow.pyhas_jammy_flows_package()is added undergraphnet.utils.importsto check if its installed, and is used in a few places to make sure that the code runs also for people who choose not to install jammy_flows.Minor changes:
repeat_labelsis added as an argument toGraphDefinition- ifTrue, event-level information, .e.genergyis repeated row-wise to match the number of pulses in the event. This feature was added in this PR to make it possible to build flows that learn pulse-level pdfs conditioned on event-level information.**kwargsforTraineris added forpredict-methods inEasySyntaxto allow the same level of control overTrainerarguments as we have for.fit