docs/helpers/device.html
1importtorch23fromlabml.configsimportBaseConfigs,hyperparams,option
6classDeviceInfo:
7def\_\_init\_\_(self,\*,8use\_cuda:bool,9cuda\_device:int):10self.use\_cuda=use\_cuda11self.cuda\_device=cuda\_device12self.cuda\_count=torch.cuda.device\_count()1314self.is\_cuda=self.use\_cudaandtorch.cuda.is\_available()15ifnotself.is\_cuda:16self.device=torch.device('cpu')17else:18ifself.cuda\_device\<self.cuda\_count:19self.device=torch.device('cuda',self.cuda\_device)20else:21self.device=torch.device('cuda',self.cuda\_count-1)
23def\_\_str\_\_(self):24ifnotself.is\_cuda:25return"CPU"2627ifself.cuda\_device\<self.cuda\_count:28returnf"GPU:{self.cuda\_device} - {torch.cuda.get\_device\_name(self.cuda\_device)}"29else:30return(f"GPU:{self.cuda\_count - 1}({self.cuda\_device}) "31f"- {torch.cuda.get\_device\_name(self.cuda\_count - 1)}")
This is a configurable module to get a single device to train model on. It can pick up CUDA devices and it will fall back to CPU if they are not available.
It has other small advantages such as being able to view the actual device name on configurations view of labml app <https://github.com/labmlai/labml/tree/master/app> _
Arguments: cuda_device (int): The CUDA device number. Defaults to 0 . use_cuda (bool): Whether to use CUDA devices. Defaults to True .
34classDeviceConfigs(BaseConfigs):
47cuda\_device:int=048use\_cuda:bool=True4950device\_info:DeviceInfo5152device:torch.device
54def\_\_init\_\_(self):55super().\_\_init\_\_(\_primary='device')
58@option(DeviceConfigs.device)59def\_device(c:DeviceConfigs):60returnc.device\_info.device616263hyperparams(DeviceConfigs.cuda\_device,DeviceConfigs.use\_cuda,64is\_hyperparam=False)656667@option(DeviceConfigs.device\_info)68def\_device\_info(c:DeviceConfigs):69returnDeviceInfo(use\_cuda=c.use\_cuda,70cuda\_device=c.cuda\_device)