search

dingo-gw / dingo

Dingo: Deep inference for gravitational-wave observations
MIT License
68 stars 22 forks source link

Unable to use Dingo Pipe with 1 detector networks due to bad size of nflows training array #216

Closed nihargupte-ph closed 5 months ago

nihargupte-ph commented 1 year ago

When doing inference with a 1 detector network, one needs to train an unconditional density estimator on the set of 1D GNPE samples. However, for some reason nflows does not like this as it thinks the training vector is the wrong shape. This occurs during the sampling step of dingo_pipe. The error message is the following: 

12:53 dingo_pipe INFO    : Command line arguments: Namespace(ini='outdir_GW190620_030421/GW190620_030421_config_complete.ini', verbose=False, calibration_model=None, spline_calibration_envelope_dict=None, spline_calibration_nodes=10, spline_calibration_curves=1000, idx=0, data_dump_file=None, event_data_files=['outdir_GW190620_030421/data/GW190620_030421_data0_1245035079-3_generation_event_data.hdf5'], proposal_samples_file=None, importance_sampling_generation=False, ignore_gwpy_data_quality_check=True, shift_segment_for_psd_generation_if_nan=True, timeslide_file=None, trigger_time='1245035079.3', data_dict=None, data_format=None, allow_tape=True, channel_dict='{H1:GWOSC, L1:GWOSC}', detectors=['L1'], duration=8.0, psd_dict=None, psd_fractional_overlap=0.0, post_trigger_duration=2.0, sampling_frequency=4096.0, psd_length=128, psd_maximum_duration=1024, psd_method='median', psd_start_time=None, maximum_frequency='1024.0', minimum_frequency='20.0', tukey_roll_off=0.4, resampling_method='lal', injection_dict=None, injection_file=None, injection_random_seed=None, asd_dataset=None, use_psd_of_trigger=True, injection_waveform_approximant=None, zero_noise=False, num_noise_realizations=100, accounting='dingo', accounting_user=None, label='GW190620_030421_data0_1245035079-3_sampling', local=False, local_generation=True, local_plot=False, outdir='/work/nihargupte/analysis/dingo/events/O3/GW190620_030421/eccentric_hm_v2/outdir_GW190620_030421', overwrite_outdir=False, request_disk=0.5, request_memory=120.0, request_memory_generation=8.0, request_cpus=1, request_cpus_importance_sampling=32, sampling_requirements=['TARGET.CUDAGlobalMemoryMb>20000'], extra_lines=['None'], simple_submission=False, conda_env=None, scheduler='condor', scheduler_args=None, scheduler_module=[None], scheduler_env=None, scheduler_analysis_time='7-00:00:00', submit=False, condor_job_priority=0, transfer_files=True, disable_hdf5_locking=True, log_directory=None, osg=False, desired_sites=None, analysis_executable=None, analysis_executable_parser=None, plot_data=False, plot_spectrogram=False, plot_corner=True, plot_weights=True, plot_log_probs=True, create_summary=False, email=None, notification='Never', queue=None, existing_dir=None, webdir=None, summarypages_arguments=None, result_format='hdf5', final_result=True, final_result_nsamples=20000, prior_dict=None, reference_frequency=None, waveform_approximant='SEOBNRv4EHM_opt', model='/data/nihargupte/projects/dingo-devel/tutorials/09_eccentricity/train_dirs/production_train_dirs/train_dir_SEOBNRv4EHM_5/model_latest.pt', model_init='/data/nihargupte/projects/dingo-devel/tutorials/09_eccentricity/train_dirs/train_dir_settings_SEOBNRv4EHM_opt_O3_L1_time/model_latest.pt', recover_log_prob=True, device='cuda', num_gnpe_iterations=30, num_samples=500000, batch_size=50000, density_recovery_settings='ProxyRecoveryDefault', importance_sample=True, importance_sampling_settings='Default', importance_sampling_updates='{}', n_parallel=20)
12:53 dingo_pipe INFO    : Unknown command line arguments: []
12:53 dingo_pipe INFO    : GNPE network does not provide log probability. Generating samples and training a new network to recover it.
/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/init.py:405: UserWarning: Initializing zero-element tensors is a no-op
  warnings.warn("Initializing zero-element tensors is a no-op")
/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/linear.py:114: UserWarning: An output with one or more elements was resized since it had shape [4096, 256], which does not match the required output shape [256]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at ../aten/src/ATen/native/Resize.cpp:26.)
  return F.linear(input, self.weight, self.bias)
Traceback (most recent call last):
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/bin/dingo_pipe_sampling", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/work/nihargupte/src/dingo-devel-aei/dingo/pipe/sampling.py", line 240, in main
    analysis.run_sampler()
  File "/work/nihargupte/src/dingo-devel-aei/dingo/pipe/sampling.py", line 183, in run_sampler
    prepare_log_prob(
  File "/work/nihargupte/src/dingo-devel-aei/dingo/gw/inference/inference_pipeline.py", line 244, in prepare_log_prob
    unconditional_model = result.train_unconditional_flow(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/src/dingo-devel-aei/dingo/core/result.py", line 468, in train_unconditional_flow
    unconditional_model = train_unconditional_density_estimator(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/src/dingo-devel-aei/dingo/core/density/unconditional_density_estimation.py", line 103, in train_unconditional_density_estimator
    model.train(
  File "/work/nihargupte/src/dingo-devel-aei/dingo/core/models/posterior_model.py", line 336, in train
    train_loss = train_epoch(self, train_loader)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/src/dingo-devel-aei/dingo/core/models/posterior_model.py", line 463, in train_epoch
    loss = -pm.model(data[0], *data[1:]).mean()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/src/dingo-devel-aei/dingo/core/nn/nsf.py", line 242, in forward
    return self.log_prob(y)
           ^^^^^^^^^^^^^^^^
  File "/work/nihargupte/src/dingo-devel-aei/dingo/core/nn/nsf.py", line 215, in log_prob
    return self.flow.log_prob(y)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/distributions/base.py", line 40, in log_prob
    return self._log_prob(inputs, context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/flows/base.py", line 39, in _log_prob
    noise, logabsdet = self._transform(inputs, context=embedded_context)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/transforms/base.py", line 56, in forward
    return self._cascade(inputs, funcs, context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/transforms/base.py", line 50, in _cascade
    outputs, logabsdet = func(outputs, context)
                         ^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/transforms/base.py", line 56, in forward
    return self._cascade(inputs, funcs, context)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/transforms/base.py", line 50, in _cascade
    outputs, logabsdet = func(outputs, context)
                         ^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/transforms/coupling.py", line 84, in forward
    transform_params = self.transform_net(identity_split, context)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/nn/nets/resnet.py", line 98, in forward
    temps = block(temps, context=context)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/glasflow/nflows/nn/nets/resnet.py", line 42, in forward
    temps = self.batch_norm_layers[0](temps)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/batchnorm.py", line 138, in forward
    self._check_input_dim(input)
  File "/work/nihargupte/environments/micromamba/envs/eccentricity5/lib/python3.11/site-packages/torch/nn/modules/batchnorm.py", line 301, in _check_input_dim
    raise ValueError(
ValueError: expected 2D or 3D input (got 1D input)
nihargupte-ph commented 1 year ago

This seems to happen because the

glasflow.nflows.transforms.coupling.PiecewiseRationalQuadraticCouplingTransform doesn't like tensors of the shape [n_batch, 1]

nihargupte-ph commented 1 year ago

I've tracked it down to the following line: https://github.com/uofgravity/nflows/blob/glasflow/nflows/transforms/coupling.py#L84

When training running this transform on a [n_batch, 1] shape tensor this identity_split is a size [n_batch, 0] tensor. Whereas if we have a 2D flow identity_split is a [n_batch, 1] size tensor.

But actually, this is a conceptual problem now that I think about it. In a coupling transform one needs at least a dimension 2 vector so that one can swap the identity portion (0th column) and the transform portion (1st column) of the vector after the transform.

A lazy workaround could be to pass a vector of 0's to the NDE as a second dummy column. I'm not sure how well this would do performance-wise since we would effectively be learning a delta-function cross the GNPE proxy distribution. (And learning a delta function for a flow is always hard). But I don't have a good intuition on this.

Otherwise, we could use a simpler set of transforms for the 1D case. That is something that doesn't use a coupling transform.

@stephengreen @jonaswildberger thoughts?