用自己的数据集训练TimeSformer并转ONNX用c++推理

用自己的数据集训练TimeSformer并转ONNX用c++推理

下载安装TimeSformer

TimeSformer开源地址

按照官方教程安装好环境。
如果报下面这个错误,是因为新版的pytorch已经不支持那种写法了，需要修改一下。

ImportError: cannot import name '_LinearWithBias' from 'torch.nn.modules.linear'

可以参考这个人的fork修改

创建分类文件夹

我这里有61个动作分类，每个分类创建一个文件夹

将视频文件分割成 每个视频大概10s左右；
然后将视频文件按照分类放到每个文件夹里。

创建数据集

写一个脚本分割数据集，并生成标签文件

import os
import csv
import shutil
from tqdm import tqdm
from sklearn.model_selection import train_test_splitout_dir = "/home/disk/liangbaikai/TimeSformer/mydata/mydatasets"  # 输出路径
video_path = "/home/disk/liangbaikai/TimeSformer/mydata/myvideos" # 数据集路径
file_name = ".csv"
name_list = ["train","test","val"]if not os.path.exists(out_dir):os.mkdir(out_dir)
if not os.path.exists(os.path.join(out_dir, 'train')):os.mkdir(os.path.join(out_dir, 'train'))
if not os.path.exists(os.path.join(out_dir, 'val')):os.mkdir(os.path.join(out_dir, 'val'))
if not os.path.exists(os.path.join(out_dir, 'test')):os.mkdir(os.path.join(out_dir, 'test'))for file in os.listdir(video_path):file_path = os.path.join(video_path, file)video_files = [name for name in os.listdir(file_path)]#将20%的数据分配给testtrain_and_valid, test = train_test_split(video_files, test_size=0.2, random_state=42)#将80%的数据再分配20%出来给val，剩下的给traintrain, val = train_test_split(train_and_valid, test_size=0.2, random_state=42)train_dir = os.path.join(out_dir, 'train', file)val_dir = os.path.join(out_dir, 'val', file)test_dir = os.path.join(out_dir, 'test', file)if not os.path.exists(train_dir):os.mkdir(train_dir)if not os.path.exists(val_dir):os.mkdir(val_dir)if not os.path.exists(test_dir):os.mkdir(test_dir)for video in tqdm(train):shutil.copy(os.path.join(video_path,file,video),os.path.join(train_dir,video))for video in tqdm(test):shutil.copy(os.path.join(video_path,file,video),os.path.join(test_dir,video))for video in tqdm(val):shutil.copy(os.path.join(video_path,file,video),os.path.join(val_dir,video))#输出路径下创建csv文件夹，并在文件夹下创建train.csv val.csv test.csv
csv_path = os.path.join(out_dir,"csv")
if not os.path.exists(csv_path):os.mkdir(csv_path)for name in name_list:with open(os.path.join(csv_path,name+file_name),'wb') as f:print("创建"+os.path.join(csv_path,name+file_name))for ii in os.listdir(csv_path):if ii.split(".")[0] in name_list:path1 = os.path.join(csv_path,ii)with open(path1, 'w', newline='') as f:for dd in os.listdir(out_dir):if dd==ii.split(".")[0]:for zz in os.listdir(os.path.join(out_dir,dd)):for mm in os.listdir(os.path.join(out_dir,dd,zz)):writer = csv.writer(f)writer.writerow([os.path.join(out_dir,dd,zz,mm),zz])## 创建类别label标号文件
labels= []
for label in sorted(os.listdir(video_path)):labels.append(label)
label2index = {label: index for index, label in enumerate(sorted(set(labels)))}
label_file = os.path.join(out_dir, str(len(os.listdir(video_path))) + 'class_labels.txt')
with open(label_file, 'w') as f:for id, label in enumerate(sorted(label2index)):f.writelines(str(id) + ' ' + label +'\n')#替换csv文件中类别名为数字
csv_file = os.path.join(out_dir,"csv")
def txt_read(files):txt_dict = {}fopen = open(files)for line in fopen.readlines():line = str(line).replace('\n','')txt_dict[line.split(' ',1)[1]] = line.split(' ',1)[0]      fopen.close()return txt_dict
txt_dict = txt_read(label_file)
print(txt_dict)for ii in os.listdir(csv_file):path1 = os.path.join(csv_file,ii)r = csv.reader(open(path1))lines = [l for l in r]for i in range(