使用Slowfast模型训练自己的数据
往期的博客介绍了视频数据、视频数据的信息获取、加载和处理，这期博客我们主要介绍Slowfast模型、写一个数据加载器，将预处理好的数据加载进Slowfast模型做训练。
Slowfast模型介绍
Slowfast模型是何凯明团队于19年的工作，不同以往的双流（光流和视频），Slowfast的双流（都是视频数据）卷积输入，一个慢通道，一个快通道，分别提取空域信息以及时域信息。Slowfast模型论文地址，模型的效果如下图：
上图的对比结果我们可以看出，Slowfast模型在19年以较低的计算量和较高的准确度取得不错的排名，我手上的一些训练任务也是使用的Slowfast模型做训练，因此我想先引入Slowfast模型，循序渐进学习更多的新模型。
编写数据加载器
数据的训练需要不断地加载数据，pytorch提供了一个数据加载接口，我们只需小小修改就可以使用了。有一篇博文讲解的很不错，博友们可以去看看：迄今为止最细致的DataSet和Dataloader加载步骤(肝)
from torch.utils.data import Dataset, DataLoader
class MyDataset(Dataset):
def __init__(self, txt, load_type="test"):
self.load_type = load_type
with open(txt, "r") as file:
self.data = file.readlines()
def __getitem__(self, idx):
video_path = self.data[idx].strip()
data = np.load(video_path)
# 数据的保存样式[20, 224*223*3 + label]
# 加载20张图片数据
imgs = data[:, :-1].reshape(20, 224, 224, 3)
label = data[:, -1:].reshape(-1)
return imgs, label[0]
def __len__(self):
return len(self.data)
加载数据做训练
我们知道数据的加载后就可以做训练了 （底下的代码可能需要有一定的知识储备，如果你有看不懂的地方可以在评论区提出，我会尽可能解惑） 。深度学习三要素：处理数据、训练数据、评估部署模型，下面的代码就是模型训练代码了。
"""
训练SlowFast模型
"""
from torch.utils.data import Dataset, DataLoader
# 加载模型
from models.model import COMP_F
import torch
import numpy as np
from tqdm import tqdm
from time import sleep
import os
torch.backends.cudnn.benchmark = True
def save_matrix(file, epoch, train_matrix, test_matrix, mode="a+"):
"""
保存混淆矩阵
"""
with open(file, mode) as f:
f.writelines(f"epoch: {epoch}\n")
f.writelines(f"train_matrix: {train_matrix}\n")
f.writelines(f"test_matrix: {test_matrix}\n")
def save_acc(file, epoch, acc, loss, lr, mode="a+"):
"""
保存训练准确度、学习率等信息
"""
with open(file, mode) as f:
f.writelines(f"{epoch}\n")
f.writelines(f"{acc}\n")
f.writelines(f"{loss}\n")
f.writelines(f"{lr}\n")
def calculate_matrix(predicted, label, matrix):
"""
计算混淆矩阵
"""
for pre_y, real_y in zip(predicted, label):
if not(f"{pre_y}_{real_y}" in matrix.keys()):
matrix[f"{pre_y}_{real_y}"] = 0
matrix[f"{pre_y}_{real_y}"] += 1
class MyDataset(Dataset):
def __init__(self, txt, load_type="test"):
self.load_type = load_type
with open(txt, "r") as file:
self.data = file.readlines()
def __getitem__(self, idx):
video_path = self.data[idx].strip()
data = np.load(video_path)
# 数据的保存样式[20, 224*223*3 + label]
# 加载20张图片数据
imgs = data[:, :-1].reshape(20, 224, 224, 3)
label = data[:, -1:].reshape(-1)
return imgs, label[0]
def __len__(self):
return len(self.data)
if __name__ == "__main__":
# 将需要训练的数据的路径写入txt
train_path = "../txt/train.txt"
test_path = "../txt/test.txt"
# 加载数据
dataset = MyDataset(train_path, load_type="train")
train_dataloader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=4, pin_memory=True)
dataset = MyDataset(test_path, load_type="test")
test_dataloader = DataLoader(dataset, batch_size=4, shuffle=False, pin_memory=True)
Epochs = 500
criterion = torch.nn.CrossEntropyLoss()
device = "cuda"
# 模型的大小
model_type = "resnet18"
os.makedirs(f"{model_type}", exist_ok=True)
os.makedirs(f"{model_type}/log", exist_ok=True)
model = COMP_F(model_type).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=1e-8)
last_acc = 0
save_matrix(file=f"{model_type}/log/matrix.txt", epoch=" ", train_matrix=" ", test_matrix=" ", mode="w")
save_acc(file=f"{model_type}/log/acc.txt", epoch=" ", acc=" ", loss=" ", lr=" ", mode="w")
for epoch in range(Epochs):
model.train()
train_loss = 0
batch_num = 0
train_correct = 0
sample_num = 0
train_matrix = {}
with open(f"{model_type}/log/异常数据.txt", "a+") as error_file:
for sample, label in tqdm(train_dataloader):
sample, label = torch.permute(sample.to(device) / 255, (0, 4, 1, 2, 3)), label.to(device)
# print(sample.shape, label.shape)
optimizer.zero_grad()
output = model(sample)
loss = criterion(output, label)
loss.backward()
optimizer.step()
batch_num += 1
train_loss += loss.item()
_, predicted = torch.max(output, 1)
sample_num += label.size(0)
train_correct += (predicted == label).sum()
# 计算混淆矩阵
calculate_matrix(predicted.cpu().numpy(), label.cpu().numpy(), train_matrix)
train_acc = train_correct.item() / sample_num
train_loss = train_loss / batch_num
# 测试
model.eval()
test_loss = 0
batch_num = 0
test_correct = 0
sample_num = 0
test_matrix = {}
with torch.no_grad():
try:
for sample, label in tqdm(test_dataloader):
sample, label = torch.permute(sample.to(device) / 255, (0, 4, 1, 2, 3)), label.to(device)
optimizer.zero_grad()
output = model(sample)
loss = criterion(output, label)
batch_num += 1
test_loss += loss.item()
_, predicted = torch.max(output, 1)
calculate_matrix(predicted.cpu().numpy(), label.cpu().numpy(), test_matrix)
sample_num += label.size(0)
test_correct += (predicted == label).sum()
except:
pass
test_acc = test_correct.item() / sample_num
torch.save(model.state_dict(), f'./{model_type}/train_{epoch}_{test_acc:.4f}.pth')
test_loss = test_loss / batch_num
save_matrix(file=f"{model_type}/log/matrix.txt", epoch=epoch, train_matrix=train_matrix, test_matrix=test_matrix)
save_acc(file=f"{model_type}/log/acc.txt", epoch=epoch, acc=f"train_acc: {train_acc:.4f} test_acc: {test_acc:.4f}"
, loss=f"training_loss:{train_loss:.5f} test_loss:{test_loss:.5f}", lr=f"Lr:{scheduler.get_last_lr()}")
print(f"Epoch: {epoch} training_loss:{train_loss:.5f} test_loss:{test_loss:.5f}")
print(f"train_acc: {train_acc:.4f} test_acc: {test_acc:.4f}")
scheduler.step()
adjusted_lr = scheduler.get_last_lr()
print(f"Epoch Lr:{adjusted_lr}\n")
sleep(1)
结尾
在这篇blog中我们了解了Slowfast模型、编写了一个数据加载器、并使用数据加载器加载数据做训练，下期我们对pytorch的开源视频模型做一下测试 （博友们有什么不懂得或者想学的也可以私信我或者在评论区写出来，我会的话就去写一写博客的）。感谢您的观看，觉得写的还可以请帮忙点个赞和收藏！
247
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
