Package: torchutil
torchutilGeneral utilities for developing deep learning projects using PyTorchpip install torchutilTable of contentsCheckpointtorchutil.checkpoint.best_pathtorchutil.checkpoint.latest_pathtorchutil.checkpoint.loadtorchutil.checkpoint.saveCudatorchutil.cuda.utilizationDownloadtorchutil.download.filetorchutil.download.tarbz2torchutil.download.targztorchutil.download.zipGradientstorchutil.gradients.statsInferencetorchutil.inference.contextIteratortorchutil.iteratortorchutil.multiprocess_iteratorMetricstorchutil.metrics.Accuracytorchutil.metrics.Averagetorchutil.metrics.F1torchutil.metrics.L1torchutil.metrics.MeanStdtorchutil.metrics.Precisiontorchutil.metrics.PearsonCorrelationtorchutil.metrics.Recalltorchutil.metrics.RMSENotifytorchutil.notifyPathstorchutil.paths.chdirtorchutil.paths.measuretorchutil.paths.purgeTensorboardtorchutil.tensorboard.updateTimetorchutil.time.contexttorchutil.time.resultstorchutil.time.resetCheckpointimporttorchimporttorchutil# Checkpoint locationfile='model.pt'# Initialize PyTorch modelmodel=torch.nn.Sequential(torch.nn.Conv1d())# Initialize optimizeroptimizer=torch.nn.Adam(model.parameters())# Savetorchutil.checkpoint.save(file,model,optimizer,step=0,epoch=0)# Load for trainingmodel,optimizer,state=torchutil.checkpoint.load(file,model,optimizer)step,epoch=state['step'],state['epoch']# Load for inferencemodel,*_=torchutil.checkpoint.load(file,model,optimizer)torchutil.checkpoint.best_pathdefbest_path(directory:Union[str,bytes,os.PathLike],glob:str='*.pt',best_fn:Callable=highest_score)->Tuple[Union[str,bytes,os.PathLike],float]:"""Retrieve the path to the best checkpointArgumentsdirectory - The directory to search for checkpoint filesglob - The regular expression matching checkpointsbest_fn - Takes a list of checkpoint paths and returns the latestDefault assumes checkpoint names are training step count.Returnsbest_file - The filename of the checkpoint with the best scorebest_score - The corresponding score"""torchutil.checkpoint.latest_pathdeflatest_path(directory:Union[str,bytes,os.PathLike],glob:str='*.pt',latest_fn:Callable=largest_number_filename,)->Union[str,bytes,os.PathLike]:"""Retrieve the path to the most recent checkpoint in a directoryArgumentsdirectory - The directory to search for checkpoint filesglob - The regular expression matching checkpointslatest_fn - Takes a list of checkpoint paths and returns the latest.Default assumes checkpoint names are training step count.ReturnsThe latest checkpoint in directory according to latest_fn"""torchutil.checkpoint.loaddefload(file:Union[str,bytes,os.PathLike],model:torch.nn.Module,optimizer:Optional[torch.optim.Optimizer]=None,map_location:str='cpu')->Tuple[torch.nn.Module,Union[None,torch.optim.Optimizer],Dict]:"""Load model checkpointArgumentsfile - The checkpoint filemodel - The PyTorch modeloptimizer - Optional PyTorch optimizer for trainingmap_location - The device to load the checkpoint onReturnsmodel - The model with restored weightsoptimizer - Optional optimizer with restored parametersstate - Additional values that the user defined during save"""torchutil.checkpoint.savedefsave(file:Union[str,bytes,os.PathLike],model:torch.nn.Module,optimizer:torch.optim.Optimizer,accelerator:Optional[accelerate.Accelerator]=None,**kwargs):"""Save training checkpoint to diskArgumentsfile - The checkpoint filemodel - The PyTorch modeloptimizer - The PyTorch optimizeraccelerator - HuggingFace Accelerator for device managementkwargs - Additional values to save"""Cudaimporttorchimporttorchutil# Directory to write Tensorboard filesdirectory='tensorboard'# Training stepstep=0# Log VRAM utilization in MB to Tensorboardtorchutil.tensorboard.update(directory,step,scalars=torchutil.cuda.utilization(torch.device('cuda:0'),'MB'))torchutil.cuda.utilizationdefutilization(device:torch.Device,unit:str='B')->Dict[str,float]:"""Get the current VRAM utilization of a specified deviceArgumentsdeviceThe device to query for VRAM utilizationunitUnit of memory utilization (bytes to terabytes); default bytesReturnsAllocated and reserved VRAM utilization in the specified unit"""Downloadtorchutil.download.filedeffile(url:'str',path:Union[str,bytes,os.PathLike]):"""Download file from urlArgumentsurl - The URL to downloadpath - The location to save results"""torchutil.download.tarbz2deftarbz2(url:'str',path:Union[str,bytes,os.PathLike]):"""Download and extract tar bz2 file to locationArgumentsurl - The URL to downloadpath - The location to save results"""torchutil.download.targzdeftargz(url:'str',path:Union[str,bytes,os.PathLike]):"""Download and extract tar gz file to locationArgumentsurl - The URL to downloadpath - The location to save results"""torchutil.download.zipdefzip(url:'str',path:Union[str,bytes,os.PathLike]):"""Download and extract zip file to locationArgumentsurl - The URL to downloadpath - The location to save results"""Gradientsimporttorchimporttorchutil# Directory to write Tensorboard filesdirectory='tensorboard'# Training stepstep=0# Setup model and optimizer# ...# Compute forward pass and loss# ...# Zero gradientsoptimizer.zero_grad()# Compute gradientsloss.backward()# Monitor gradients on tensorboardtorchutil.tensorboard.update(directory,step,scalars=torchutil.gradients.stats(model))# Apply gradient updateoptimizer.step()torchutil.gradients.statsdefstats(model:torch.nn.Module)->Dict[str,float]:"""Get gradient statisticsArgumentsmodelThe torch modelReturnsThe L2 norm, maximum, and minimum gradients"""Inferencetorchutil.inference.context@contextlib.contextmanagerdefcontext(model:torch.nn.Module,autocast:bool=True)->None:"""Inference-time handling of model training flag and optimizationsArgumentsmodelThe torch model performing inferenceautocastWhether to use mixed precision"""Iteratorimporttimeimporttorchutildefwait(seconds):time.sleep(seconds)n=8iterable=range(n)# Monitor single-process jobforiintorchutil.iterator(iterable,message='single-process'):wait(i)# Monitor multi-process jobtorchutil.multiprocess_iterator(wait,iterable,message='multi-process')torchutil.iteratordefiterator(iterable:Iterable,message:Optional[str]=None,initial:int=0,total:Optional[int]=None)->Iterable:"""Create a tqdm iteratorArgumentsiterableItems to iterate overmessageStatic message to displayinitialPosition to display corresponding to index zero of iterabletotalLength of the iterable; defaults to len(iterable)ReturnsMonitored iterable"""torchutil.multiprocess_iteratordefmultiprocess_iterator(process:Callable,iterable:Iterable,message:Optional[str]=None,initial:int=0,total:Optional[int]=None,num_workers:int=os.cpu_count(),worker_chunk_size:Optional[int]=None)->List:"""Create a multiprocess tqdm iteratorArgumentsprocessThe single-argument function called by each multiprocess workeriterableItems to iterate overmessageStatic message to displayinitialPosition to display corresponding to index zero of iterabletotalLength of the iterable; defaults to len(iterable)num_workersMultiprocessing pool size; defaults to number of logical CPU coresworker_chunk_sizeNumber of items sent to each multiprocessing workerReturnsReturn values of calling process on each item, in original order"""Metricsimporttorchimporttorchutil# Define a custom, batch-updating loss metricclassLoss(torchutil.metrics.Average):defupdate(self,predicted,target):# Compute your loss and the number of elements to average overloss=...count=...super().update(loss,count)# Instantiate metricsloss=Loss()rmse=torchutil.metrics.RMSE()# Generator that produces batches of predicted and target tensorsiterable=...# Update metricsforpredicted_tensor,target_tensoriniterable:loss.update(predicted_tensor,target_tensor)rmse.update(predicted_tensor,target_tensor)# Get resultsprint('loss':loss())print('rmse':rmse())torchutil.metrics.AccuracyclassAccuracy(Metric):"""Batch-updating accuracy metric"""def__call__(self)->float:"""Retrieve the current accuracy valueReturnsThe current accuracy value"""defupdate(self,predicted:torch.Tensor,target:torch.Tensor)->None:"""Update accuracyArgumentspredictedThe model predictiontargetThe corresponding ground truth"""defreset(self)->None:"""Reset accuracy"""torchutil.metrics.AverageclassAverage(Metric):"""Batch-updating average metric"""def__call__(self)->float:"""Retrieve the current average valueReturnsThe current average value"""defupdate(self,values:torch.Tensor,count:int)->None:"""Update running averageArgumentsvaluesThe values to averagetargetThe number of values"""defreset(self)->None:"""Reset running average"""torchutil.metrics.F1classF1(Metric):"""Batch-updating F1 score"""def__call__(self)->float:"""Retrieve the current F1 valueReturnsThe current F1 value"""defupdate(self,predicted:torch.Tensor,target:torch.Tensor)->None:"""Update F1ArgumentspredictedThe model predictiontargetThe corresponding ground truth"""defreset(self)->None:"""Reset F1"""torchutil.metrics.L1classL1(Metric):"""Batch updating L1 score"""def__call__(self)->float:"""Retrieve the current L1 valueReturnsThe current L1 value"""defupdate(self,predicted,target)->None:"""Update L1ArgumentspredictedThe model predictiontargetThe corresponding ground truth"""defreset(self)->None:"""Reset L1"""torchutil.metrics.MeanStdclassMeanStd(Metric):"""Batch updating mean and standard deviation"""def__call__(self)->Tuple[float,float]:"""Retrieve the current mean and standard deviationReturnsThe current mean and standard deviation"""defupdate(self,values:torch.Tensor)->None:"""Update mean and standard deviationArgumentspredictedThe model predictiontargetThe corresponding ground truth"""defreset(self)->None:"""Reset mean and standard deviation"""torchutil.metrics.PearsonCorrelationclassPearsonCorrelation(Metric):"""Batch-updating Pearson correlation"""def__init__(self,predicted_mean:float,predicted_std:float,target_mean:float,target_std:float)->None:"""Argumentspredicted_mean - Mean of predicted valuespredicted_std - Standard deviation of predicted valuestarget_mean - Mean of target valuestarget_std - Standard deviation of target values"""def__call__(self)->float:"""Retrieve the current correlation valueReturnsThe current correlation value"""defupdate(self,predicted,target)->None:"""Update Pearson correlationArgumentspredictedThe model predictiontargetThe corresponding ground truth"""defreset(self)->None:"""Reset Pearson correlation"""torchutil.metrics.PrecisionclassPrecision(Metric):"""Batch-updating precision metric"""def__call__(self)->float:"""Retrieve the current precision valueReturnsThe current precision value"""defupdate(self,predicted:torch.Tensor,target:torch.Tensor)->None:"""Update precisionArgumentspredictedThe model predictiontargetThe corresponding ground truth"""defreset(self)->None:"""Reset precision"""torchutil.metrics.RecallclassRecall(Metric):"""Batch-updating recall metric"""def__call__(self)->float:"""Retrieve the current recall valueThe current recall value"""defupdate(self,predicted:torch.Tensor,target:torch.Tensor)->None:"""Update recallArgumentspredictedThe model predictiontargetThe corresponding ground truth"""defreset(self)->None:"""Reset recall"""torchutil.metrics.RMSEclassRMSE(Metric):"""Batch-updating RMSE metric"""def__call__(self)->float:"""Retrieve the current rmse valueReturnsThe current rmse value"""defupdate(self,predicted:torch.Tensor,target:torch.Tensor)->None:"""Update RMSEArgumentspredictedThe model predictiontargetThe corresponding ground truth"""defreset(self)->None:"""Reset RMSE"""NotifyTo use thetorchutilnotification system, set thePYTORCH_NOTIFICATION_URLenvironment variable to a supported webhook as explained inthe Apprise documentation.importtorchutil# Send notification when function returns@torchutil.notify('train')deftrain():...# Equivalent using context managerdeftrain():withtorchutil.notify('train'):...torchutil.notify@contextlib.contextmanagerdefnotify(description:str,track_time:bool=True,notify_on_fail:bool=True):"""Context manager for sending job notificationsArgumentsdescription - The name of the job being runtrack_time - Whether to report time elapsednotify_on_fail - Whether to send a notification on failure"""Pathstorchutil.paths.chdir@contextlib.contextmanagerdefchdir(directory:Union[str,bytes,os.PathLike])->None:"""Context manager for changing the current working directoryArgumentsdirectoryThe desired working directory"""This function is both a context manager and decorator.importtempfilefrompathlibimportPathimporttorchutil# Create a directorydirectory=tempfile.TemporaryDirectory()# Create a filefile='tmp.txt'(Path(directory.name)/file).touch()# File is not in current working directoryassertnotPath(file).exists()# Change working directory using context managerwithtorchutil.paths.chdir(directory.name):assertPath(file).exists()# File is not in current working directoryassertnotPath(file).exists()# Change working directory using decorator@torchutil.paths.chdir(directory.name)defexists(file):assertPath(file).exists()exists(file)# File is not in current working directoryassertnotPath(file).exists()# Remove temporary pathsdirectory.cleanup()torchutil.paths.measuredefmeasure(globs:Optional[List[Union[str,List[str]]]]=None,roots:Optional[List[Union[Union[str,bytes,os.PathLike],List[Union[str,bytes,os.PathLike]]]]]=None,recursive:bool=False,unit='B')->Union[int,float]:"""Measure data usage of files and directoriesArgumentsglobsGlobs matching paths to measurerootsDirectories to apply glob searches; current directory by defaultrecursiveApply globs to all subdirectories of root directoriesunitUnit of memory utilization (bytes to terabytes); default bytesReturnsData usage in the specified unit"""This function also has a command-line interface.python -m torchutil.paths.measure \
    [-h] \
    --globs GLOBS \
    [--roots ROOTS] \
    [--recursive] \
    [--unit]

Measure data usage of files and directories

arguments:
  --globs GLOBS
    Globs matching paths to measure

optional arguments:
  -h, --help
    show this help message and exit
  --roots ROOTS
    Directories to apply glob searches; current directory by default
  --recursive
    Apply globs to all subdirectories of root directories
  --unit
    Unit of memory utilization (bytes to terabytes); default bytestorchutil.paths.purgedefpurge(globs:Optional[List[Union[str,List[str]]]]=None,roots:Optional[List[Union[Union[str,bytes,os.PathLike],List[Union[str,bytes,os.PathLike]]]]]=None,recursive:bool=False,force:bool=False)->None:"""Remove all files and directories within directory matching globArgumentsglobsGlobs matching paths to deleterootsDirectories to apply glob searches; current directory by defaultrecursiveApply globs to all subdirectories of root directoriesforceSkip user confirmation of deletion"""This function also has a command-line interface.python -m torchutil.paths.purge \
    [-h] \
    --globs GLOBS \
    [--roots ROOTS] \
    [--recursive] \
    [--force]

Remove files and directories

arguments:
  --globs GLOBS
    Globs matching paths to delete

optional arguments:
  -h, --help
    show this help message and exit
  --roots ROOTS
    Directories to apply glob searches; current directory by default
  --recursive
    Apply globs to all subdirectories of root directories
  --force
    Skip user confirmation of deletionTensorboardimportmatplotlib.pyplotaspltimporttorchimporttorchutil# Directory to write Tensorboard filesdirectory='tensorboard'# Training stepstep=0# Example audioaudio=torch.zeros(1,16000)sample_rate=16000# Example figurefigure=plt.figure()plt.plot([0,1,2,3])# Example imageimage=torch.zeros(256,256,3)# Example scalarloss=0# Update Tensorboardtorchutil.tensorboard.update(directory,step,audio={'audio':audio},sample_rate=sample_rate,figures={'figure':figure},images={'image':image},scalars={'loss':loss})torchutil.tensorboard.updatedefupdate(directory:Union[str,bytes,os.PathLike],step:int,audio:Optional[Dict[str,torch.Tensor]]=None,sample_rate:Optional[int]=None,figures:Optional[Dict[str,matplotlib.figure.Figure]]=None,images:Optional[Dict[str,torch.Tensor]]=None,scalars:Optional[Dict[str,Union[float,int,torch.Tensor]]]=None):"""Update TensorboardArgumentsdirectory - Directory to write Tensorboard filesstep - Training stepaudio - Optional dictionary of 2D audio tensors to monitorsample_rate - Audio sample rate; required if audio is not Nonefigures - Optional dictionary of Matplotlib figures to monitorimages - Optional dictionary of 3D image tensors to monitorscalars - Optional dictionary of scalars to monitor"""Timeimporttimeimporttorchutil# Perform timingwithtorchutil.time.context('outer'):time.sleep(1)foriinrange(2):time.sleep(1)withtorchutil.time.context('inner'):time.sleep(1)# Prints {'inner': 2.0020763874053955, 'outer': 5.005248308181763, 'total': 5.005248308181763}print(torchutil.time.results())torchutil.time.context@contextlib.contextmanagerdefcontext(name:str):"""Wrapper to handle context changes of global timerArgumentsname - Name of the timer to add time to"""torchutil.time.resetdefreset():"""Clear timer state"""torchutil.time.resultsdefresults()->dict:"""Get timing resultsReturnsTiming results: {name: elapsed_time} for all names"""