NTU60RGBD数据集训练速度慢问题和解决

一直向前669

于 2024-04-23 15:13:09 发布

阅读量333

点赞数 5

文章标签： python 计算机视觉人工智能机器学习

本文链接：https://blog.csdn.net/qq_60925525/article/details/138125352

版权

数据集深度图展示

创作原因

本人最近在进行计算机视觉的研究，用到了该数据集，在运行该数据集的过程中遇到了问题，训练速度极慢，cpu和gpu使用率低,io活跃，最终解决了问题，在查找资料的过程中，发现网络上基本没有对于相关问题的讨论，借此机会多谈谈我的想法。

问题原因

本人的实验设备为4块3090显卡，个人认为可以称得上是非常高配了。

数据集大小，单纯的ntu60深度数据原始数据约为88G,ntu120更大,超过了100G，本文仅讨论ntu60，其实这里大小还行，但是对于数据预处理后大小来到了665G，我的个人PC已经装不下了，只能自掏腰包买了个移动硬盘（1T）。

数据集加载代码如下：

import os
import sys
import numpy as np
from torch.utils.data import Dataset

Cross_Subject = [1, 2, 4, 5, 8, 9, 13, 14, 15, 16, 17, 18, 19, 25, 27, 28, 31, 34, 35, 38]

def clip_normalize(clip):
    pc = np.reshape(a=clip, newshape=[-1, 3])
    centroid = np.mean(pc, axis=0)
    pc = pc - centroid
    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
    clip = (clip - centroid) / m
    return clip

class CLRNTU60Subject(Dataset):
    def __init__(self, root, meta, frames_per_clip=23, step_between_clips=2, num_points=2048, step_between_frames=2, train=True):
        super(CLRNTU60Subject, self).__init__()

        self.videos = []
        self.index_map = []
        index = 0

        with open(meta, 'r') as f:
            for line in f:
                name, nframes = line.split()
                subject = int(name[9:12])
                if train:
                    if subject in Cross_Subject:
                        label = int(name[-3:]) - 1
                        nframes = int(nframes)
                        for t in range(0, nframes-step_between_frames*(frames_per_clip-1), step_between_clips):
                            self.index_map.append((index, t))
                        index += 1
                        self.videos.append(os.path.join(root, name+'.npz'))

        self.frames_per_clip = frames_per_clip
        self.step_between_clips = step_between_clips
        self.step_between_frames = step_between_frames
        self.num_points = num_points
        self.train = train


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

    def __getitem__(self, idx):
        index, t = self.index_map[idx]

        video = self.videos[index]
        video = np.load(video, allow_pickle=True)['data'] * 100

        clip = [video[t+i*self.step_between_frames] for i in range(self.frames_per_clip)]
        for i, p in enumerate(clip):
            if p.shape[0] > self.num_points:
                r = np.random.choice(p.shape[0], size=self.num_points, replace=False)
            else:
                repeat, residue = self.num_points // p.shape[0], self.num_points % p.shape[0]
                r = np.random.choice(p.shape[0], size=residue, replace=False)
                r = np.concatenate([np.arange(p.shape[0]) for _ in range(repeat)] + [r], axis=0)
            clip[i] = p[r, :]
        clip = clip_normalize(np.array(clip))

        if self.train:
            # scale the points
            scales = np.random.uniform(0.9, 1.1, size=3)
            clip = clip * scales

        clip = clip.astype(np.float32) # [L, N, 3]

        return clip, index

Dataset中，init方法仅加载路径，直到调用getitem才会从磁盘中读取对应数据。

665G数据无法直接装入内存，因此需要把数据先装入先读入内存再送给gpu运行，而由于磁盘的访问速度较慢，这里我们就发现磁盘的io速度成为了瓶颈。个人测试过，对于较小数据集可以直接装入内存，就不存在该问题，即使不在init将所有数据装入，速度也不会受到太大影响，因为大部分数据在运行第一轮后还在内存中，并不会被移出内存，后续运行时可以直接在内存中找到，减少了io次数。