util.py 文件源码

def get_torch_num_workers(num_workers: int):
    """turn an int into a useful number of workers for a pytorch DataLoader.
    -1 means "use all CPU's", -2, means "use all but 1 CPU", etc.
    Note: 0 is interpreted by pytorch as doing data loading in the main process, while any positive number spawns a
    new process. We do not allow more processes to spawn than there are CPU's."""
    num_cpu = cpu_count()
    if num_workers < 0:
        n_workers = num_cpu + 1 + num_workers
        if n_workers < 0:
            print("Warning: {} fewer workers than the number of CPU's were specified, but there are only {} CPU's; "
                  "running data loading in the main process (num_workers = 0).".format(num_workers + 1, num_cpu))
        num_workers = max(0, n_workers)
    if num_workers > num_cpu:
        print("Warning, `num_workers` is {} but only {} CPU's are available; "
              "using this number instead".format(num_workers, num_cpu))
    return min(num_workers, num_cpu)


#####################################################################
# Iterator utils                                                    #
#####################################################################