gputools.py 文件源码

def __init__(self):
        self.devices = device_lib.list_local_devices()
        self.cpus = [x.name for x in self.devices if x.device_type == 'CPU']
        self.gpus = [x.name for x in self.devices if x.device_type == 'GPU']
        self.iterate_cpus = it.cycle(self.cpus)
        self.iterate_gpus = it.cycle(self.gpus)