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

本文的目的是为了与相同研究方向的读者进行交流和学习。这也是本人最近在动作识别领域复现的第一篇论文，所以以这样的方式记录下来。本人也坚信，不管做啥事情都要自己去努力，不然别人也帮不了自己。整篇文章的内容主要讲的是本人复现代码的记录，所以需要了解双流网络原理的朋友，去参考学习原论文（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)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

六、训练可视化结果

图2 训练过程可视化

七、致谢和说明

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

八、数据和源码地址

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