动作识别(一)复现双流网络

本文的目的是为了与相同研究方向的读者进行交流和学习。这也是本人最近在动作识别领域复现的第一篇论文，所以以这样的方式记录下来。本人也坚信，不管做啥事情都要自己去努力，不然别人也帮不了自己。整篇文章的内容主要讲的是本人复现代码的记录，所以需要了解双流网络原理的朋友，去参考学习原论文（Two-Stream Convolutional Networks for Action Recognition in Videos），网上链接很多。

图1 用于视频分类的双流架构

一、本代码运行环境

(1)win10;
(2)cuda10.1+RTX2060显卡,自查；
(3)conda 创建环境: conda create -n pytorch python=3.6.2;
(4)安装pyrtorch: pip install torch1.7.1+cu101 torchvision0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
(5)安装其他库: pip install tensorboard, tensorboardX, tqdm, matplotlib

二、代码目录解释

|-data:存放所有数据
|----RGB：存放RGB图像
|----UCF101：存放原始视频数据，需要自己找链接下载，也可联系本人2447439418@qq.com
|-Gnerate_RGB&FLOW:生成需要的数据，文件内有说明，本人认为很清楚
|-model:存放模型和日志文件
|----models:日志文件
|----.pth:训练好的模型
|-ucfTrainTestlist:数据划分文件
|-其余文件，打开顶部都有解释

三、可视化日志文件步骤

(1)打开终端，切换当前环境和代码目录;
(2)输入tensorboard --logdir model 命令在网页查看训练过程

四、注意事项

(1)RGB和Flow网络的主干网络均为Resent101,RGB网络输入大小为(10, 3, 224, 224),输出为(10, 2);Flow网络输入大小为(10, 20, 224, 224),输出为(10, 2)；其中RGB网络采用预训练模型，故开始运行代码时要连接网络，从Pytorch官网下载与训练模型，否则报错；
(2)光流图像可划分为x和y方向上的两个图像，故Flow网络的输入中堆叠光流为20,也即堆叠了10帧光流图像
(3)生成RGB+Flow图像的代码在Gnerate_RGB&FLOW文件中，文档中有说明；
(4)初次下载测试模型和处理好的数据，直接替换当前代码文件中的model和data即可

五、Two-stream训练网络部分代码

"""
训练文件
"""
from LoadUCF101Data import trainset_loader, testset_loader
from two_Stream_Net import TwoStreamNet
import torch
import torch.optim as optim
import os
import socket
import torch.nn.functional as F
from tqdm import tqdm
from datetime import datetime
from tensorboardX import SummaryWriter


# 超参数
EPOCH = 35  # 总训练轮数
LEARNING_RATE = 1e-3  # 初始学习率
MOMENTUM = 0.9  # 动量因子
N_TEST = 2  # 每N_TEST轮测试一次
N_SAVE = 5  # 每N_SAVE轮保存一次模型
SAVE_DIR = "./model"
# 选择设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 双流网络模型
twoStreamNet = TwoStreamNet().to(device)
# SGD优化器
optimizer = optim.SGD(
    params=twoStreamNet.parameters(),
    lr=LEARNING_RATE,
    momentum=MOMENTUM,
)
# 学习率衰减
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)  # the scheduler divides the lr by 10 every 10 epochs

# 保存模型
def save_checkpoint(path, model, optimizer):
    state = {
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict()
    }
    torch.save(state, path)


# 测试函数
def test(i_epoch, epoch, writer):
    twoStreamNet.eval()

    runing_correct = 0
    runing_loss = 0
    with torch.no_grad():
        for data in tqdm(testset_loader):
            # 加载batch_size个数据
            RGB_images, OpticalFlow_images, label = data
            RGB_images = RGB_images.to(device)
            OpticalFlow_images = OpticalFlow_images.to(device)
            label = label.to(device)
            # 计算batch_size损失
            output = twoStreamNet(RGB_images, OpticalFlow_images)
            loss = F.cross_entropy(output, label)
            runing_loss += loss.item() * len(label)
            # 计算测试集batch_size中正确的个数
            max_value, max_index = output.max(1, keepdim=True)
            runing_correct += max_index.eq(label.view_as(max_index)).sum().item()
        # 计算每个epoch上的平均损失
        epoch_loss = runing_loss / len(testset_loader.dataset)
        epoch_acc = (runing_correct * 1.0 * 100 / len(testset_loader.dataset))  # 计算准确度

        writer.add_scalar('data/test_loss_epoch', epoch_loss, i_epoch-1)
        writer.add_scalar('data/test_acc_epoch', epoch_acc, i_epoch-1)

        print("[{}] Epoch: {}/{} Loss: {} Acc: {}".format('test', i_epoch, epoch, epoch_loss, epoch_acc))

# 训练函数
def train(epoch, N_TEST, N_SAVE, save_dir):
    # 存入日志文件，方便使用tensorboard工具观看训练变化
    log_dir = os.path.join(save_dir, 'models', datetime.now().strftime('%b%d_%H-%M-%S') + '_' + socket.gethostname())
    writer = SummaryWriter(log_dir=log_dir)
    for i in range(epoch):
        twoStreamNet.train()

        runing_correct = 0
        runing_loss = 0
        for data in tqdm(trainset_loader):
            # 加载bach_size个数据
            RGB_images, OpticalFlow_images, label = data
            RGB_images = RGB_images.to(device)
            OpticalFlow_images = OpticalFlow_images.to(device)
            label = label.to(device)
            # 更新一次模型参数
            optimizer.zero_grad()
            output = twoStreamNet(RGB_images, OpticalFlow_images)
            loss = F.cross_entropy(output, label)
            loss.backward()
            optimizer.step()
            # 计算当前训练损失
            runing_loss += loss.item() * len(label)
            # 计算当前正确个数
            max_value, max_index = output.max(1, keepdim=True)
            runing_correct += max_index.eq(label.view_as(max_index)).sum().item()

        scheduler.step()  # 学习率衰减策略
        print("第%d个epoch的学习率：%f" % (i + 1, optimizer.param_groups[0]['lr']))

        # 计算每个epoch上的平均损失
        epoch_loss = runing_loss / len(trainset_loader.dataset)
        epoch_acc = (runing_correct * 1.0 * 100 / len(trainset_loader.dataset))  # 计算准确度
        # 写入日志文件中
        writer.add_scalar('data/train_loss_epoch', epoch_loss, i)
        writer.add_scalar('data/train_acc_epoch', epoch_acc, i)
        # 打印当前轮结果
        print("[{}] Epoch: {}/{} Loss: {} Acc: {}".format('train', i + 1, epoch, epoch_loss, epoch_acc))
        # 训练N_TEST轮测试一次
        if (i+1) % N_TEST == 0:
            test(i + 1, epoch, writer)
        # 训练N_SAVE轮测试一次，保存一次模型
        if (i+1) % N_SAVE == 0:
            save = save_dir + "/checkpoint-%i.pth"
            save_checkpoint(save % (i+1), twoStreamNet, optimizer)
            print("Save checkpoint-{}.pth at {}\n".format(i+1, "./model"))



if __name__ == '__main__':
    train(EPOCH, N_TEST, N_SAVE, SAVE_DIR)

六、训练可视化结果

图2 训练过程可视化

七、致谢和说明

本代码参考了大量的作者代码，非常感谢它们！本人代码也要开源！也给自己写一点笔记，嘿嘿！下一篇博客名字已经想好了，【动作识别（二）复现C3D网络】，有时间一定记录下来，感谢大家的交流！

八、数据和源码地址

强调一下，由于UCF101原始数据太大了，本人百度云空间有限，所以将UCF101中的前两个数据拿出来了，并生成了相应的RGB+Flow数据!
处理好的数据地址.提取码：fj2j
测试模型地址.提取码：mp86
DMK-Two-stream源码.提取码：g643