判断服务器的可用cpu数量来设置num_workers

在一个多人使用的集群环境中，直接使用 multiprocessing.cpu_count() 获取的 CPU 核心数量可能不准确，因为部分 CPU 核心可能已经被其他用户的任务占用。为了更准确地确定可用的 CPU 核心数量，你可以使用一些系统工具来检查当前系统的 CPU 负载情况。

使用 psutil 检查可用 CPU 核心

psutil 是一个跨平台的库，可以轻松获取系统的运行信息，包括 CPU 负载、内存使用情况等。下面是一个示例代码，展示如何使用 psutil 来获取可用 CPU 核心数量：

import psutil

def get_available_cpu_count():
    """
    获取当前系统的可用CPU核心数量
    :return: 可用的CPU核心数量
    """
    # 获取每个CPU核心的当前负载
    cpu_percentages = psutil.cpu_percent(percpu=True)
    
    # 设置一个阈值，例如CPU负载小于50%时认为是可用的
    threshold = 50.0
    available_cores = sum(cpu < threshold for cpu in cpu_percentages)
    
    return available_cores

available_cpu_count = get_available_cpu_count()
print(f"Available CPU cores: {available_cpu_count}")

使用 psutil 和 multiprocessing 综合判断

为了更全面地获取系统的 CPU 信息，可以结合 psutil 和 multiprocessing 来判断总的 CPU 核心数量和当前可用的 CPU 核心数量。

import psutil
import multiprocessing

def get_total_and_available_cpu_count():
    """
    获取总的CPU核心数量和当前系统的可用CPU核心数量
    :return: 总的CPU核心数量, 可用的CPU核心数量
    """
    total_cpu_count = multiprocessing.cpu_count()
    cpu_percentages = psutil.cpu_percent(percpu=True)
    
    # 设置一个阈值，例如CPU负载小于50%时认为是可用的
    threshold = 50.0
    available_cores = sum(cpu < threshold for cpu in cpu_percentages)
    
    return total_cpu_count, available_cores

total_cpu_count, available_cpu_count = get_total_and_available_cpu_count()
print(f"Total CPU cores: {total_cpu_count}")
print(f"Available CPU cores: {available_cpu_count}")

示例代码：根据可用 CPU 核心数设置 num_workers

在获取了可用的 CPU 核心数量后，你可以根据这个数量来设置 num_workers。

import os
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import psutil
import multiprocessing

class InMemoryVideoFramesDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.dataset_dirs = [os.path.join(root_dir, d) for d in os.listdir(root_dir) if os.path.isdir(os.path.join(root_dir, d))]
        self.frames_list = self.load_all_frames_to_memory()
    
    def load_all_frames_to_memory(self):
        all_frames = []
        for dataset_dir in self.dataset_dirs:
            frames_list = sorted([os.path.join(dataset_dir, f) for f in os.listdir(dataset_dir) if f.endswith('.raw')])
            frames = [self.read_raw_file(frame_path) for frame_path in frames_list]
            all_frames.append(frames)
        return all_frames
    
    def __len__(self):
        return len(self.frames_list)
    
    def __getitem__(self, idx):
        frames = self.frames_list[idx]
        if self.transform:
            frames = [self.transform(frame) for frame in frames]
        return torch.stack(frames, dim=0)
    
    @staticmethod
    def read_raw_file(file_path):
        with open(file_path, 'rb') as f:
            frame = np.fromfile(f, dtype=np.uint8).reshape((480, 640))
        return torch.from_numpy(frame).float() / 255.0

def get_available_cpu_count():
    cpu_percentages = psutil.cpu_percent(percpu=True)
    threshold = 50.0
    available_cores = sum(cpu < threshold for cpu in cpu_percentages)
    return available_cores

# 数据集根目录
root_dir = 'path_to_your_root_directory'

# 创建数据集
dataset = InMemoryVideoFramesDataset(root_dir)

# 获取可用的CPU核心数
available_cpu_count = get_available_cpu_count()
num_workers = available_cpu_count // 2

# 创建数据加载器
dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=num_workers)

# 测试数据加载器
for batch in dataloader:
    print(batch.shape)

总结

通过使用 psutil 获取当前系统的 CPU 负载信息，可以更准确地确定可用的 CPU 核心数量，并根据这个数量来设置 num_workers。这样可以更好地利用系统资源，提高数据加载和训练的效率。如果你有进一步的问题或需要更多帮助，请随时告诉我。