GPUAccelerator¶

class pytorch_lightning.accelerators.GPUAccelerator(precision_plugin, training_type_plugin)[source]¶

Accelerator for GPU devices.

Parameters

precision_plugin¶ (PrecisionPlugin) – the plugin to handle precision-specific parts
training_type_plugin¶ (TrainingTypePlugin) – the plugin to handle different training routines

static auto_device_count()[source]¶

Get the devices when set to auto.

get_device_stats(device)[source]¶

Gets stats for the given GPU device.

Parameters: device¶ (Union[str, device]) – GPU device for which to get stats
Return type: Dict[str, Any]
Returns: A dictionary mapping the metrics to their values.
Raises: FileNotFoundError – If nvidia-smi installation not found

on_train_start()[source]¶

Called when train begins.

setup(trainer)[source]¶

Setup plugins for the trainer fit and creates optimizers.

setup_environment()[source]¶

This method is called to teardown the training process.

It is the right place to release memory and free other resources.