Added verbose arg stump to datashift and spikedetect

Author: jacobpennington

kilosort/clustering_qr.py

@@ -502,6 +502,9 @@ def run(ops, st, tF,  mode = 'template', device=torch.device('cuda'),
 
                 logger.debug(f'Center {ii} | Xd shape: {Xd.shape} | ntemp: {ntemp}')
                 if verbose and Xd.nelement() > 10**8:
+                    logger.info(f'Resetting cuda memory stats for Center {ii}')
+                    if device == torch.device('cuda'):
+                        torch.cuda.reset_peak_memory_stats(device)
                     v = True
 
                 if Xd is None:

kilosort/datashift.py

@@ -184,7 +184,7 @@ def kernel2D(x, y, sig = 1):
     return Kn
 
 def run(ops, bfile, device=torch.device('cuda'), progress_bar=None,
-        clear_cache=False):
+        clear_cache=False, verbose=False):
     """ this step computes a drift correction model
     it returns vertical correction amplitudes for each batch, and for multiple blocks in a batch if nblocks > 1. 
     """
@@ -197,7 +197,7 @@ def run(ops, bfile, device=torch.device('cuda'), progress_bar=None,
     # the first step is to extract all spikes using the universal templates
     st, _, ops  = spikedetect.run(
         ops, bfile, device=device, progress_bar=progress_bar,
-        clear_cache=clear_cache
+        clear_cache=clear_cache, verbose=verbose
         )
 
     # spikes are binned by amplitude and y-position to construct a "fingerprint" for each batch

kilosort/run_kilosort.py

@@ -232,6 +232,7 @@ def run_kilosort(settings, probe=None, probe_name=None, filename=None,
         ops, bfile, st0 = compute_drift_correction(
             ops, device, tic0=tic0, progress_bar=progress_bar,
             file_object=file_object, clear_cache=clear_cache,
+            verbose=verbose_log
             )
 
         # Check scale of data for log file
@@ -531,7 +532,7 @@ def compute_preprocessing(ops, device, tic0=np.nan, file_object=None):
 
 
 def compute_drift_correction(ops, device, tic0=np.nan, progress_bar=None,
-                             file_object=None, clear_cache=False):
+                             file_object=None, clear_cache=False, verbose=False):
     """Compute drift correction parameters and save them to `ops`.
 
     Parameters
@@ -548,6 +549,11 @@ def compute_drift_correction(ops, device, tic0=np.nan, progress_bar=None,
         Must have 'shape' and 'dtype' attributes and support array-like
         indexing (e.g. [:100,:], [5, 7:10], etc). For example, a numpy
         array or memmap.
+    clear_cache : bool; False.
+        If True, force pytorch to clear cached cuda memory after some
+        memory-intensive steps in the pipeline.
+    verbose : bool; False.
+        If true, include additional debug-level logging statements.
 
     Returns
     -------
@@ -580,7 +586,7 @@ def compute_drift_correction(ops, device, tic0=np.nan, progress_bar=None,
         )
 
     ops, st = datashift.run(ops, bfile, device=device, progress_bar=progress_bar,
-                            clear_cache=clear_cache)
+                            clear_cache=clear_cache, verbose=verbose)
     bfile.close()
     logger.info(f'drift computed in {time.time()-tic : .2f}s; ' + 
                 f'total {time.time()-tic0 : .2f}s')
@@ -650,7 +656,7 @@ def detect_spikes(ops, device, bfile, tic0=np.nan, progress_bar=None,
     logger.info('-'*40)
     st0, tF, ops = spikedetect.run(
         ops, bfile, device=device, progress_bar=progress_bar,
-        clear_cache=clear_cache
+        clear_cache=clear_cache, verbose=verbose
         )
     tF = torch.from_numpy(tF)
     logger.info(f'{len(st0)} spikes extracted in {time.time()-tic : .2f}s; ' + 

kilosort/spikedetect.py

@@ -190,7 +190,7 @@ def yweighted(yc, iC, adist, xy, device=torch.device('cuda')):
     return yct
 
 def run(ops, bfile, device=torch.device('cuda'), progress_bar=None,
-        clear_cache=False):        
+        clear_cache=False, verbose=False):        
     sig = ops['settings']['min_template_size']
     nsizes = ops['settings']['template_sizes']