DMOPR & MLOPR currently has no GPU acceleration

[ltang320]

Hi Benedikt @daurer ,
I want to process a dataset with probe changing during the data collection. I am using multi-modes with DMOPR and MLOPR algorithms.

It is extremely slow. Does it support GPU to accelerate? Or do you have some good suggestions to make it faster?

  
  ---- Probe initialization ------------------------------------------------------
  Initializing probe storage Sscan00G00 using scan scan00.
  Found no photon count for probe in parameters.
  Using photon count 9.43e+06 from photon report
  WARNING ptypy - You are doing a multimodal reconstruction with none/ not much diversity between the modes! 
  This will likely not reconstruct. You should set .scan.illumination.diversity.power and .scan.illumination.diversity.noise to something for the best results.
  [Object Sscan00G00] Model illumination is propagated over a distance 0.012 m.
  
  ---- Object initialization -----------------------------------------------------
  Initializing object storage Sscan00G00 using scan scan00.
  [Object Sscan00G00] Simulation resource is integrated refractive index
  
  ---- Creating exit waves -------------------------------------------------------
  
  Process #0 ---- Total Pods 10200 (10200 active) ----
  --------------------------------------------------------------------------------
  (C)ontnr : Memory : Shape            : Pixel size      : Dimensions      : Views
  (S)torgs : (MB)   : (Pixel)          : (meters)        : (meters)        : act. 
  --------------------------------------------------------------------------------
  Cprobe   : 5347.7 : complex64
  Sscan00G00 : 5347.7 : 0200 * 256 * 256 :      5.0*5.0e-8 :      1.3*1.3e-5 : 10200
  Cobj     :   12.4 : complex64
  Sscan00G00 :   12.4 :  1 * 1248 * 1238 :      5.0*5.0e-8 :      6.3*6.2e-5 : 10200
  Cexit    : 5347.7 : complex64
  S0000G00 : 5347.7 : 0200 * 256 * 256 :      5.0*5.0e-8 :      1.3*1.3e-5 : 10200
  Cdiff    :  668.5 : float32
  S0000    :  668.5 : 2550 * 256 * 256 :      7.5*7.5e-5 :      1.9*1.9e-2 :  2550
  Cmask    :  167.1 :   bool
  S0000    :  167.1 : 2550 * 256 * 256 :      7.5*7.5e-5 :      1.9*1.9e-2 :  2550
  
  
  
  ---- Ptycho init level 3 -------------------------------------------------------
  
  ---- Ptycho init level 4 -------------------------------------------------------
  
  ==== Starting DMOPR-algorithm. =================================================
  
  Parameter set:
  * id1BC3QJ8E7G           : ptypy.utils.parameters.Param(25)
    * numiter              : 500
    * numiter_contiguous   : 5
    * probe_support        : 0.7
    * probe_fourier_sup... : None
    * record_local_error   : False
    * position_refinement  : ptypy.utils.parameters.Param(10)
      * method             : Annealing
      * start              : 100
      * stop               : None
      * interval           : 50
      * nshifts            : 256
      * amplitude          : 1e-07
      * amplitude_decay    : True
      * max_shift          : 5e-07
      * metric             : fourier
      * record             : True
    * probe_update_start   : 2
    * subpix_start         : 0
    * subpix               : linear
    * update_object_first  : True
    * overlap_converge_... : 0.05
    * overlap_max_itera... : 2
    * probe_inertia        : 1e-09
    * object_inertia       : 0.0001
    * fourier_power_bound  : None
    * fourier_relax_factor : 0.05
    * obj_smooth_std       : None
    * clip_object          : None
    * probe_center_tol     : None
    * compute_log_likel... : True
    * alpha                : 1.0
    * name                 : DMOPR
    * IP_metric            : 1.0
    * subspace_dim         : 5
    * subspace_start       : 2
  ================================================================================
  Initialising position refinement (Annealing)
  ---------------------------------- Autosaving ----------------------------------
  WARNING ptypy - Save file exists but will be overwritten (force_overwrite is True)
  Generating copies of probe, object and parameters and runtime
  Saving to /asap3/petra3/gpfs/p06/2023/data/11017776/processed/alignment/scan_00062/ptypy_LT_single_GPU_crop_256_testOPR/dumps/dump_scan_00062_None_0000.ptyr
   --------------------------------------------------------------------------------
       Time spent in Fourier update: 62.98
  Time spent in Overlap update: 1419.57
  Time spent in Position update: 0.00
  Iteration #5 of DMOPR :: Time 1482.555
  Errors :: Fourier 1.01e+02, Photons 1.32e+02, Exit 1.24e+02

[daurer]

@ltang320 there is currently no GPU version of the OPR engines but you can use MPI to speed it up. For example, if you have access to an HPC cluster with N CPUs, you could run your DMOPR or MLOPR script with

mpirun -n <N> python run_ptypy_dmopr.py