基于pytorch实现任意python任务多线程处理

不啰嗦，直接上代码！

from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader

__all__ = ["MultiProcessor", "taskCore"]


class taskCore(Dataset):
    def __init__(self, custom_list, *args, **kwargs):
        self.custom_list = custom_list
        self.args = kwargs
        self.kwargs = kwargs
        pass

    def __len__(self):
        return len(self.custom_list)

    def pro_func(self, *args, **kwargs):
        """替换for循环处理内容"""
        pass

    def custom_item(self, custom_list, idx):
        """替换从for循环中取数据"""
        return custom_list[idx]

    def __getitem__(self, idx):
        return self.pro_func(self.custom_item(self.custom_list, idx), *self.args, **self.kwargs)


def collate_fn(item):
    return 0


class MultiProcessor:
    def __init__(self, custom_list, custom_taskCore=None, *args, **kwargs):
        self.taskCore = custom_taskCore if custom_taskCore else taskCore(custom_list, *args, **kwargs)
        self.handler = DataLoader(dataset=self.taskCore, collate_fn=collate_fn, num_workers=8)

    def set_custom_func(self, pro_func, custom_item=None):
        self.taskCore.pro_func = pro_func
        if custom_item is not None:
            self.taskCore.custom_item = custom_item

    def __call__(self, *args, **kwargs):
        for item in tqdm(self.handler):
            pass


def main():
    pass


if __name__ == '__main__':
    main()
    pass