6.6 PWC-Net算法
PWC-Net（Pyramidal Warping-Net）使用金字塔结构来处理不同尺度的信息，通过金字塔池化和金字塔卷积融合多尺度特征，提高了光流估计的准确性。请看下面的实例，基于PWC-Net光流估计技术实现了一个视频字幕生成系统，本项目由如下功能模块构成：
- 特征提取模块：使用 PWC-Net 模型提取视频的光流特征，使用 I3D 模型提取视频的 RGB 特征。
- 字幕提议生成模块：利用预训练的提议生成模型（prop_model）结合音频、RGB 和光流特征生成初始的字幕提议。
- 非极大值抑制：可选的步骤，对生成的字幕提议进行非极大值抑制。
- 字幕生成模块：利用预训练的字幕生成模型（cap_model）结合音频、RGB、光流特征以及生成的字幕提议为每个提议生成最终的字幕描述。
整个项目通过上述一系列的模块协同工作，可以从给定的视频中生成具有语义描述的字幕，提高了视频内容的理解和表达。同时，项目中还包括了一些附加功能，如支持并行处理多个视频、提供了一些可选的输出方式等。
实例6-8：视频字幕生成系统（codes/6/project-bmt.ipynb）
实例文件project-bmt.ipynb的具体实现代码如下所示。
6.6.1 准备环境和数据
（1）通过下面的命令下载预训练模型和数据文件，然后将它们移动到项目中的指定目录。
glove.840B.300d.zip 100%[===================>] 2.03G 16.9MB/s in 2m 8s
best_cap_model.pt 100%[===================>] 589.61M 17.3MB/s in 36s
best_prop_model.pt 100%[===================>] 2.50G 17.4MB/s in 2m 32s
vggish_model.ckpt 100%[===================>] 277.62M 188MB/s in 1.5s对上述命令的具体说明如下：
- 首先，使用wget命令从指定的URL下载了几个文件，包括GloVe词向量文件、字幕模型和属性模型的预训练权重文件以及VGGish模型的检查点文件。通过使用-q选项，代码在下载过程中不显示详细信息，而使用--show-progress选项则显示下载进度。
- 然后，创建了一个名为.vector_cache的目录，这是为了存储GloVe词向量文件。此外，将下载的文件移动到项目中的不同目录，以确保它们在后续的使用中能够方便地被访问到。具体来说，GloVe词向量文件被移动到.vector_cache目录，字幕模型和属性模型的权重文件被移动到./sample/目录，而VGGish模型的检查点文件则被移动到./submodules/video_features/models/vggish/checkpoints/目录。
- 最后，所有这些操作旨在准备项目所需的预训练文件和模型，以便后续步骤中的特征提取、字幕生成等任务能够顺利进行。这种结构化的文件准备过程有助于确保项目的依赖项和数据文件都被正确安装和定位。
执行后输出：
glove.840B.300d.zip 100%[===================>] 2.03G 16.9MB/s in 2m 8s
best_cap_model.pt 100%[===================>] 589.61M 17.3MB/s in 36s
best_prop_model.pt 100%[===================>] 2.50G 17.4MB/s in 2m 32s
vggish_model.ckpt 100%[===================>] 277.62M 188MB/s in 1.5s（2）导入本项目所需要的Python库，在众多的导入库中，有如下所示的库需要重点说明：
- from datasets.captioning_dataset import ActivityNetCaptionsDataset：导入自定义模块，其中包含一个用于处理视频字幕数据的数据集类。
- from datasets.load_features import crop_a_segment, pad_segment：导入自定义模块，其中包含一些处理视频特征加载的功能，如裁剪和填充。
- from epoch_loops.captioning_epoch_loops import make_masks, greedy_decoder：导入自定义模块，其中包含用于训练和推理时的循环操作的功能，如创建掩码和贪婪解码。
- from model.captioning_module import BiModalTransformer：导入自定义模块，其中包含一个双模态（图像和文本）转换器模型。
- from model.proposal_generator import MultimodalProposalGenerator：导入自定义模块，其中包含一个多模态提议生成器模型。
- from utilities.proposal_utils import ...：导入一系列自定义功能，用于处理提议的一些实用工具，如获取坐标、非最大抑制等。
- from sample.single_video_prediction import get_video_duration：导入自定义功能，用于获取视频的时长。
- from typing import Dict, List, Union：导入类型提示的相关模块，用于给函数参数和返回值添加类型提示，提高代码的可读性和维护性。
6.6.2 视频分析工具集
下面的代码实现了一个视频分析的工具集，其中 Config 类存储配置信息，包括加载特征、提议生成模型和字幕生成模型。load_features_from_npy 函数用于加载视频特征，而 generate_proposals 函数通过预训练提议生成模型生成视频提议。最后，caption_proposals 函数使用预训练字幕生成模型为提议生成字幕。这些工具可用于分析视频内容并生成字幕描述。
（1）定义函数 load_features_from_npy()，用于从numpy文件中加载预提取的视频特征。
def load_features_from_npy(
feature_paths: Dict[str, str], start: float, end: float, duration: float, pad_idx: int,
device: int, get_full_feat=False, pad_feats_up_to: Dict[str, int] = None
) -> Dict[str, torch.Tensor]:
'''从numpy文件中加载预提取的特征。
该函数在概念上类似于`datasets.load_feature.load_features_from_npy`，但经过整理以进行演示。
参数:
feature_paths (Dict[str, str]): numpy文件的路径（键：'audio'，'rgb'，'flow'）。
start (float, None): 提议的开始时间（秒），如果用于字幕生成提议。
end (float, None): 提议的结束时间（秒），如果用于字幕生成提议。
duration (float): 原始视频的持续时间（秒）。
pad_idx (int): 训练词汇中填充标记的索引。
device (int): GPU id。
get_full_feat (bool, optional): 是否输出完整、未修剪的特征堆栈。默认为False。
pad_feats_up_to (Dict[str, int], optional): 如果get_full_feat，填充到此值。对于音频和视频模态不同。默认为None。
返回:
Dict[str, torch.Tensor]: 包含'audio'，'rgb'和'flow'特征的字典。
'''
# 加载特征。请查看根文件夹中的README以获取有关视频特征提取的信息
stack_vggish = np.load(feature_paths['audio'])
stack_rgb = np.load(feature_paths['rgb'])
stack_flow = np.load(feature_paths['flow'])
stack_vggish = torch.from_numpy(stack_vggish).float()
stack_rgb = torch.from_numpy(stack_rgb).float()
stack_flow = torch.from_numpy(stack_flow).float()
# 用于提议生成，我们对特征进行填充
if get_full_feat:
stack_vggish = pad_segment(stack_vggish, pad_feats_up_to['audio'], pad_idx)
stack_rgb = pad_segment(stack_rgb, pad_feats_up_to['video'], pad_idx)
stack_flow = pad_segment(stack_flow, pad_feats_up_to['video'], pad_idx=0)
# 对于字幕生成，裁剪与提议相对应的片段
else:
stack_vggish = crop_a_segment(stack_vggish, start, end, duration)
stack_rgb = crop_a_segment(stack_rgb, start, end, duration)
stack_flow = crop_a_segment(stack_flow, start, end, duration)
# 添加批次维度，发送到设备
stack_vggish = stack_vggish.to(torch.device(device)).unsqueeze(0)
stack_rgb = stack_rgb.to(torch.device(device)).unsqueeze(0)
stack_flow = stack_flow.to(torch.device(device)).unsqueeze(0)
return {'audio': stack_vggish, 'rgb': stack_rgb, 'flow': stack_flow}上述代码的实现流程如下所示：
- 首先，通过参数feature_paths中的键（'audio'，'rgb'，'flow'）加载预提取的音频、RGB和光流特征。这些特征被存储在对应的numpy文件中。
- 然后，这些加载的特征通过NumPy的load函数转换为NumPy数组，并接着被转换为PyTorch张量，具体分别对应stack_vggish、stack_rgb和stack_flow。
接着，根据参数get_full_feat的取值，分两种情况处理特征：
- 如果get_full_feat为True，表示需要输出完整的、未修剪的特征堆栈。在这种情况下，函数调用自定义的pad_segment函数，将音频和视频特征堆栈分别进行填充，确保它们的长度达到指定的填充长度。
- 如果get_full_feat为False，表示需要对特定提议对应的片段进行修剪。在这种情况下，函数调用自定义的crop_a_segment函数，将音频和视频特征堆栈分别裁剪为与提议对应的时间段。
- 最后，将处理后的音频、RGB和光流特征堆栈添加批次维度，并将其发送到指定的GPU设备。最终，函数返回一个包含 'audio'，'rgb' 和 'flow' 特征的字典，这个字典将作为输入传递给模型。
（2）定义函数load_prop_model()，用于加载预训练的提议生成模型和用于训练模型的配置对象。
def load_prop_model(
device: int, prop_generator_model_path: str, pretrained_cap_model_path: str, max_prop_per_vid: int
) -> tuple:
'''加载预训练的提议生成器和用于训练模型的配置对象。
Args:
device (int): GPU id。
prop_generator_model_path (str): 预训练的提议生成模型的路径。
pretrained_cap_model_path (str): 预训练的字幕模块的路径（提议生成器使用编码器权重）。
max_prop_per_vid (int): 每个视频的最大提议数。
Returns:
Config, torch.nn.Module: 配置对象，提议生成器
'''
# 首先，加载并调整配置以适应用户定义的参数
checkpoint = torch.load(prop_generator_model_path, map_location='cpu')
cfg = checkpoint['config']
cfg.device = device
cfg.max_prop_per_vid = max_prop_per_vid
cfg.pretrained_cap_model_path = pretrained_cap_model_path
cfg.train_meta_path = './data/train.csv' # 在保存的配置中它的名字不同
# 然后，加载锚点信息
anchors = {
'audio': checkpoint['anchors']['audio'],
'video': checkpoint['anchors']['video']
}
# 接着，定义模型并加载权重
model = MultimodalProposalGenerator(cfg, anchors)
device = torch.device(cfg.device)
torch.cuda.set_device(device)
model.load_state_dict(checkpoint['model_state_dict']) # 如果有IncompatibleKeys - 忽略
model = model.to(cfg.device)
model.eval()
# 最后，返回配置对象和提议生成器模型
return cfg, model上述代码的实现流程如下；
- 首先，加载预训练的提议生成器模型的检查点。通过torch.load函数，它从指定路径prop_generator_model_path中加载了模型的权重、配置和其他相关信息。
- 然后，函数创建了一个配置对象cfg，并根据用户提供的参数对其进行调整。其中，cfg.device被设置为指定的GPU设备，cfg.max_prop_per_vid为每个视频设置最大的提议数量，cfg.pretrained_cap_model_path为预训练的字幕模块的路径，cfg.train_meta_path为训练元数据的路径。
- 接着，函数加载了锚点信息，这些信息在提议生成过程中起到重要作用。锚点用于定义提议的形状和大小，分别对应音频和视频模态。
- 然后，函数定义了提议生成器模型，并将其初始化为MultimodalProposalGenerator类的实例。这一步骤使用配置对象cfg和加载的锚点信息。接着，函数将模型移动到指定的GPU设备，并加载了预训练的权重。
- 最后，函数将模型设置为评估模式（model.eval()），以确保在推理时不进行梯度更新，并返回配置对象cfg和提议生成器模型。
（3）定义函数load_cap_model()，用于加载预训练的字幕生成模型、用于训练模型的配置对象和训练数据集，返回配置对象、字幕生成模型和训练数据集。
def load_cap_model(pretrained_cap_model_path: str, device: int) -> tuple:
cap_model_cpt = torch.load(pretrained_cap_model_path, map_location='cpu')
cfg = cap_model_cpt['config']
cfg.device = device
cfg.pretrained_cap_model_path = pretrained_cap_model_path
cfg.train_meta_path = './data/train.csv'
train_dataset = ActivityNetCaptionsDataset(cfg, 'train', get_full_feat=False)
model = BiModalTransformer(cfg, train_dataset)
model = torch.nn.DataParallel(model, [device])
model.load_state_dict(cap_model_cpt['model_state_dict']) # if IncompatibleKeys - ignore
model.eval()
return cfg, model, train_dataset上述代码都是些流程如下所示：
- 首先，函数通过torch.load加载了预训练的字幕生成模型的检查点，得到了包含配置和模型权重的字典。
- 然后，创建了一个配置对象cfg，并根据用户提供的参数进行调整。其中，cfg.device被设置为指定的GPU设备，cfg.pretrained_cap_model_path为预训练字幕模块的路径，cfg.train_meta_path为训练元数据的路径。
- 接着，函数加载了训练数据集，但仅用于获取特殊标记的索引。这一步骤通过实例化ActivityNetCaptionsDataset类来实现，其中使用了配置对象cfg，数据集的模式被设置为训练模式，且不获取完整的特征。
- 然后，函数定义了字幕生成模型，并加载了预训练的权重。模型被初始化为BiModalTransformer类的实例，并通过torch.nn.DataParallel包裹，以便在多个GPU上进行并行计算。
- 最后，函数将模型设置为评估模式（model.eval()），确保在推理时不进行梯度更新，并返回了配置对象cfg、字幕生成模型和用于构建词汇表的训练数据集。
（4）定义函数generate_proposals()，功能是使用预训练的提议生成模型生成提议，返回一个包含预测提议的张量。
# 加载特征
feature_stacks = load_features_from_npy(
feature_paths, None, None, duration_in_secs, pad_idx, device, get_full_feat=True,
pad_feats_up_to=cfg.pad_feats_up_to
)
# 形成输入批次
batch = {
'feature_stacks': feature_stacks,
'duration_in_secs': duration_in_secs
}
with torch.no_grad():
# 在输入特征中掩盖填充
masks = make_masks(batch['feature_stacks'], None, cfg.modality, pad_idx)
# 推断调用
predictions, _, _, _ = prop_model(batch['feature_stacks'], None, masks)
# （中心，长度）->（开始，结束）
predictions = get_corner_coords(predictions)
# 保持开始和结束点的合理范围
predictions = trim_proposals(predictions, batch['duration_in_secs'])
# 过滤掉长度为0或太短（<0.2）的段，使其成为提议的候选
predictions = remove_very_short_segments(predictions, shortest_segment_prior=0.2)
# 选择前 k 个预测作为最终提议
predictions = select_topk_predictions(predictions, k=cfg.max_prop_per_vid)（5）定义函数caption_proposals()，功能是使用预训练的字幕生成模型对提议进行字幕生成，返回一个包含每个提议的开始时间、结束时间和字幕的字典列表。
def caption_proposals(
cap_model: torch.nn.Module, feature_paths: Dict[str, str],
train_dataset: torch.utils.data.dataset.Dataset, cfg: Config, device: int, proposals: torch.Tensor,
duration_in_secs: float
) -> List[Dict[str, Union[float, str]]]:
'''使用预训练模型为提议生成字幕。必须指定原始视频的持续时间。
Args:
cap_model (torch.nn.Module): 预训练的字幕模型。使用load_cap_model()函数获取。
feature_paths (Dict[str, str]): 包含特征路径的字典（'audio'、'rgb'和'flow'）。
train_dataset (torch.utils.data.dataset.Dataset): 用作词汇表和特殊标记的训练数据集。
cfg (Config): 用于训练字幕模型的配置对象。预训练模型检查点包含它。
device (int): 用于计算的 GPU ID。
proposals (torch.Tensor): 大小为（batch=1，num_props，3）的张量，其中包含预测的提议。
duration_in_secs (float): 视频的持续时间（秒）。可以使用以下工具获取持续时间：
`ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 in.mp4`
Returns:
List(Dict(str, Union(float, str))): 一个字典列表，其中键为 'start'、'end' 和 'sentence'。
'''
results = []
with torch.no_grad():
for start, end, conf in proposals.squeeze():
# 加载特征
feature_stacks = load_features_from_npy(
feature_paths, start, end, duration_in_secs, train_dataset.pad_idx, device
)
# 逐个字幕单词为每个段落解码
ints_stack = greedy_decoder(
cap_model, feature_stacks, cfg.max_len, train_dataset.start_idx, train_dataset.end_idx,
train_dataset.pad_idx, cfg.modality
)
assert len(ints_stack) == 1, '该函数已清理，只支持 batch=1（validation_1by1_loop）'
# 将整数转换为字符串
strings = [train_dataset.train_vocab.itos[i] for i in ints_stack[0].cpu().numpy()]
# 删除起始标记
strings = strings[1:]
# 并删除结束标记后的所有内容
# 有时它不在列表中（当字幕意图大于 cfg.max_len 时）
try:
first_entry_of_eos = strings.index('</s>')
strings = strings[:first_entry_of_eos]
except ValueError:
pass
# 将所有内容连接在一起
sentence = ' '.join(strings)
# 将句子首字母大写
sentence = sentence.capitalize()
# 将结果添加到列表中
results.append({
'start': round(start.item(), 1),
'end': round(end.item(), 1),
'sentence': sentence
})
return results（6）定义函数which_ffprobe()，用于确定库ffprobe的路径。
def which_ffprobe() -> str:
result = subprocess.run(['which', 'ffprobe'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
ffprobe_path = result.stdout.decode('utf-8').replace('\n', '')
return ffprobe_path（7）定义函数get_video_duration()，用于确定自定义视频的持续时间，返回视频的持续时间（秒）。
def get_video_duration(path):
cmd = f'{which_ffprobe()} -hide_banner -loglevel panic' \
f' -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 {path}'
result = subprocess.run(cmd.split(), stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
video_duration = float(result.stdout.decode('utf-8').replace('\n', ''))
print('Video Duration:', video_duration)
return video_duration6.6.3 计算相关性算法
"相关性"（correlation）指的是图像特征之间的相互关系，在本项目中，相关性的计算涉及使用输入的两个张量（tensorFirst 和 tensorSecond）来生成输出张量，这个过程被称为相关性计算。相关性的计算通常涉及对两个输入张量的滑动窗口操作，计算它们在空间上的相似性。这可以在不同的视觉任务中发挥作用，例如目标检测、光流估计等。
（1）定义CUDA 核函数kernel_Correlation_rearrange，用于实现相关性计算中的重新排列操作。这个操作是为了将输入特征张量进行重新排列，以便在相关性计算中进行有效的操作。具体来说，这个 CUDA 核函数包含了一个名为 kernel_Correlation_rearrange 的函数，该函数在 GPU 上执行。它通过使用线程块和线程索引来处理输入张量中的元素，进行相关性计算所需的重新排列。
import torch
import cupy
import re
class Stream:
if torch.cuda.is_available():
ptr = torch.cuda.current_stream().cuda_stream
else:
ptr = None
# end
kernel_Correlation_rearrange = '''
extern "C" __global__ void kernel_Correlation_rearrange(
const int n,
const float* input,
float* output
) {
int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x;
if (intIndex >= n) {
return;
}
int intSample = blockIdx.z;
int intChannel = blockIdx.y;
float dblValue = input[(((intSample * SIZE_1(input)) + intChannel) * SIZE_2(input) * SIZE_3(input)) + intIndex];
__syncthreads();
int intPaddedY = (intIndex / SIZE_3(input)) + 4;
int intPaddedX = (intIndex % SIZE_3(input)) + 4;
int intRearrange = ((SIZE_3(input) + 8) * intPaddedY) + intPaddedX;
output[(((intSample * SIZE_1(output) * SIZE_2(output)) + intRearrange) * SIZE_1(input)) + intChannel] = dblValue;
}
'''（2）定义CUDA 核函数kernel_Correlation_updateOutput，用于计算相关性的输出。该函数在 GPU 上执行，通过处理输入张量 rbot0 和 rbot1 的元素，计算相关性并将结果写入输出张量 top 中。
kernel_Correlation_updateOutput = '''
extern "C" __global__ void kernel_Correlation_updateOutput(
const int n,
const float* rbot0,
const float* rbot1,
float* top
) {
extern __shared__ char patch_data_char[];
float *patch_data = (float *)patch_data_char;
// First (upper left) position of kernel upper-left corner in current center position of neighborhood in image 1
int x1 = blockIdx.x + 4;
int y1 = blockIdx.y + 4;
int item = blockIdx.z;
int ch_off = threadIdx.x;
// Load 3D patch into shared shared memory
for (int j = 0; j < 1; j++) { // HEIGHT
for (int i = 0; i < 1; i++) { // WIDTH
int ji_off = (j + i) * SIZE_3(rbot0);
for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
int idx1 = ((item * SIZE_1(rbot0) + y1+j) * SIZE_2(rbot0) + x1+i) * SIZE_3(rbot0) + ch;
int idxPatchData = ji_off + ch;
patch_data[idxPatchData] = rbot0[idx1];
}
}
}
__syncthreads();
__shared__ float sum[32];
// Compute correlation
for (int top_channel = 0; top_channel < SIZE_1(top); top_channel++) {
sum[ch_off] = 0;
int s2o = top_channel % 9 - 4;
int s2p = top_channel / 9 - 4;
for (int j = 0; j < 1; j++) { // HEIGHT
for (int i = 0; i < 1; i++) { // WIDTH
int ji_off = (j + i) * SIZE_3(rbot0);
for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
int x2 = x1 + s2o;
int y2 = y1 + s2p;
int idxPatchData = ji_off + ch;
int idx2 = ((item * SIZE_1(rbot0) + y2+j) * SIZE_2(rbot0) + x2+i) * SIZE_3(rbot0) + ch;
sum[ch_off] += patch_data[idxPatchData] * rbot1[idx2];
}
}
}
__syncthreads();
if (ch_off == 0) {
float total_sum = 0;
for (int idx = 0; idx < 32; idx++) {
total_sum += sum[idx];
}
const int sumelems = SIZE_3(rbot0);
const int index = ((top_channel*SIZE_2(top) + blockIdx.y)*SIZE_3(top))+blockIdx.x;
top[index + item*SIZE_1(top)*SIZE_2(top)*SIZE_3(top)] = total_sum / (float)sumelems;
}
}
}
'''上述代码的实现流程如下：
- 将输入张量 rbot0 的某个局部区域加载到共享内存中，以便后续计算使用。这个局部区域的大小为 1x1，通过线程块和线程索引进行加载。
- 在加载的数据上进行相关性计算，结果存储在 sum 数组中。这里使用了一个大小为 32 的数组来保存计算的中间结果。
- 将计算的相关性结果进行规约，最终在输出张量 top 中写入最终的相关性值。
值得注意的是，这个 CUDA 核函数使用了共享内存来存储中间计算结果，以提高计算效率。此外，为了避免线程同步问题，使用了 __syncthreads() 函数来进行同步。
注意：这段代码是 PWC-Net (Pyramid, Warping, and Cost Volume) 中的一部分，用于光流估计。CUDA 核函数的编写涉及到并行计算和内存访问模式的优化，以提高计算性能。
（3）定义CUDA 核函数kernel_Correlation_updateGradFirst，用于计算相关性操作的梯度。该函数通过处理输入张量 rbot0、rbot1 和 gradOutput 的元素，计算相关性操作对 rbot0 的梯度，并将结果写入 gradFirst 张量中。
kernel_Correlation_updateGradFirst = '''
#define ROUND_OFF 50000
extern "C" __global__ void kernel_Correlation_updateGradFirst(
const int n,
const int intSample,
const float* rbot0,
const float* rbot1,
const float* gradOutput,
float* gradFirst,
float* gradSecond
) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
int n = intIndex % SIZE_1(gradFirst); // channels
int l = (intIndex / SIZE_1(gradFirst)) % SIZE_3(gradFirst) + 4; // w-pos
int m = (intIndex / SIZE_1(gradFirst) / SIZE_3(gradFirst)) % SIZE_2(gradFirst) + 4; // h-pos
// round_off is a trick to enable integer division with ceil, even for negative numbers
// We use a large offset, for the inner part not to become negative.
const int round_off = ROUND_OFF;
const int round_off_s1 = round_off;
// We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
int xmin = (l - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
int ymin = (m - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
// Same here:
int xmax = (l - 4 + round_off_s1) - round_off; // floor (l - 4)
int ymax = (m - 4 + round_off_s1) - round_off; // floor (m - 4)
float sum = 0;
if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
xmin = max(0,xmin);
xmax = min(SIZE_3(gradOutput)-1,xmax);
ymin = max(0,ymin);
ymax = min(SIZE_2(gradOutput)-1,ymax);
for (int p = -4; p <= 4; p++) {
for (int o = -4; o <= 4; o++) {
// Get rbot1 data:
int s2o = o;
int s2p = p;
int idxbot1 = ((intSample * SIZE_1(rbot0) + (m+s2p)) * SIZE_2(rbot0) + (l+s2o)) * SIZE_3(rbot0) + n;
float bot1tmp = rbot1[idxbot1]; // rbot1[l+s2o,m+s2p,n]
// Index offset for gradOutput in following loops:
int op = (p+4) * 9 + (o+4); // index[o,p]
int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
for (int y = ymin; y <= ymax; y++) {
for (int x = xmin; x <= xmax; x++) {
int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
sum += gradOutput[idxgradOutput] * bot1tmp;
}
}
}
}
}
const int sumelems = SIZE_1(gradFirst);
const int bot0index = ((n * SIZE_2(gradFirst)) + (m-4)) * SIZE_3(gradFirst) + (l-4);
gradFirst[bot0index + intSample*SIZE_1(gradFirst)*SIZE_2(gradFirst)*SIZE_3(gradFirst)] = sum / (float)sumelems;
} }
'''（4）定义CUDA 核函数kernel_Correlation_updateGradSecond，用于计算相关性操作的第二个输入 rbot1 的梯度。该函数通过处理输入张量 rbot0、rbot1 和 gradOutput 的元素，计算相关性操作对 rbot1 的梯度，并将结果写入 gradSecond 张量中。
kernel_Correlation_updateGradSecond = '''
# 定义 ROUND_OFF 为 50000
extern "C" __global__ void kernel_Correlation_updateGradSecond(
const int n,
const int intSample,
const float* rbot0,
const float* rbot1,
const float* gradOutput,
float* gradFirst,
float* gradSecond
) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
int n = intIndex % SIZE_1(gradSecond); // 通道数
int l = (intIndex / SIZE_1(gradSecond)) % SIZE_3(gradSecond) + 4; // w-pos
int m = (intIndex / SIZE_1(gradSecond) / SIZE_3(gradSecond)) % SIZE_2(gradSecond) + 4; // h-pos
// round_off 是一种使整数除法可以取整数的技巧，即使是负数也可以
// 我们使用一个大的偏移量，以防内部变成负数。
const int round_off = ROUND_OFF;
const int round_off_s1 = round_off;
float sum = 0;
for (int p = -4; p <= 4; p++) {
for (int o = -4; o <= 4; o++) {
int s2o = o;
int s2p = p;
// 获取X,Y范围并进行截断
// 我们在整数除法之前添加 round_off，之后减去 round_off，以确保公式与 ceil 行为匹配：
int xmin = (l - 4 - s2o + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
int ymin = (m - 4 - s2p + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
// 同样在这里：
int xmax = (l - 4 - s2o + round_off_s1) - round_off; // floor (l - 4 - s2o)
int ymax = (m - 4 - s2p + round_off_s1) - round_off; // floor (m - 4 - s2p)
if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
xmin = max(0,xmin);
xmax = min(SIZE_3(gradOutput)-1,xmax);
ymin = max(0,ymin);
ymax = min(SIZE_2(gradOutput)-1,ymax);
// 获取 rbot0 数据：
int idxbot0 = ((intSample * SIZE_1(rbot0) + (m-s2p)) * SIZE_2(rbot0) + (l-s2o)) * SIZE_3(rbot0) + n;
float bot0tmp = rbot0[idxbot0]; // rbot1[l+s2o,m+s2p,n]
// 在以下循环中的 gradOutput 的索引偏移量：
int op = (p+4) * 9 + (o+4); // index[o,p]
int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
for (int y = ymin; y <= ymax; y++) {
for (int x = xmin; x <= xmax; x++) {
int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
sum += gradOutput[idxgradOutput] * bot0tmp;
}
}
}
}
}
const int sumelems = SIZE_1(gradSecond);
const int bot1index = ((n * SIZE_2(gradSecond)) + (m-4)) * SIZE_3(gradSecond) + (l-4);
gradSecond[bot1index + intSample*SIZE_1(gradSecond)*SIZE_2(gradSecond)*SIZE_3(gradSecond)] = sum / (float)sumelems;
}
}（5）定义函数cupy_kernel，这个函数接受两个参数：strFunction 和 objectVariables。该函数主要用于生成 CUDA C++ 的核函数字符串，其中包含一些用于处理张量维度和索引的逻辑。函数cupy_kernel的目的是根据给定的张量大小和索引信息，生成相应的 CUDA 核函数字符串，以便后续在 GPU 上执行。这是一种在运行时根据张量信息动态生成 CUDA 核函数的方法。
def cupy_kernel(strFunction, objectVariables):
strKernel = globals()[strFunction]
while True:
objectMatch = re.search('(SIZE_)([0-4])(\()([^\)]*)(\))', strKernel)
if objectMatch is None:
break
intArg = int(objectMatch.group(2))
strTensor = objectMatch.group(4)
intSizes = objectVariables[strTensor].size()
strKernel = strKernel.replace(objectMatch.group(), str(intSizes[intArg]))
while True:
objectMatch = re.search('(VALUE_)([0-4])(\()([^\)]+)(\))', strKernel)
if objectMatch is None:
break
# end
intArgs = int(objectMatch.group(2))
strArgs = objectMatch.group(4).split(',')
strTensor = strArgs[0]
intStrides = objectVariables[strTensor].stride()
strIndex = [ '((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(intStrides[intArg]) + ')' for intArg in range(intArgs) ]
strKernel = strKernel.replace(objectMatch.group(0), strTensor + '[' + str.join('+', strIndex) + ']')
return strKernel上述代码的实现流程如下：
- 从全局变量中获取与给定 strFunction 相对应的 CUDA 核函数字符串 strKernel。
- 使用正则表达式查找 strKernel 中类似于 SIZE_ 和 VALUE_ 的模式，然后替换为相应的张量维度大小和索引计算表达式。
- 遍历所有匹配项，替换 SIZE_ 模式为对应张量的大小，替换 VALUE_ 模式为对应张量的索引计算表达式。
- 返回更新后的 CUDA 核函数字符串。
（6）定义PyTorch 自动求导函数（autograd function）_FunctionCorrelation ，用于实现相关性计算的前向传播逻辑。这是一个典型的 PyTorch 自动求导函数的前向传播实现，其中涉及到 GPU 计算和 CUDA 核函数的调用。
class _FunctionCorrelation(torch.autograd.Function):
@staticmethod
def forward(self, first, second):
rbot0 = first.new_zeros([ first.size(0), first.size(2) + 8, first.size(3) + 8, first.size(1) ])
rbot1 = first.new_zeros([ first.size(0), first.size(2) + 8, first.size(3) + 8, first.size(1) ])
self.save_for_backward(first, second, rbot0, rbot1)
assert(first.is_contiguous() == True)
assert(second.is_contiguous() == True)
output = first.new_zeros([ first.size(0), 81, first.size(2), first.size(3) ])
if first.is_cuda == True:
n = first.size(2) * first.size(3)
cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
'input': first,
'output': rbot0
}))(
grid=tuple([ int((n + 16 - 1) / 16), first.size(1), first.size(0) ]),
block=tuple([ 16, 1, 1 ]),
args=[ n, first.data_ptr(), rbot0.data_ptr() ],
stream=Stream
)
n = second.size(2) * second.size(3)
cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
'input': second,
'output': rbot1
}))(
grid=tuple([ int((n + 16 - 1) / 16), second.size(1), second.size(0) ]),
block=tuple([ 16, 1, 1 ]),
args=[ n, second.data_ptr(), rbot1.data_ptr() ],
stream=Stream
)
n = output.size(1) * output.size(2) * output.size(3)
cupy_launch('kernel_Correlation_updateOutput', cupy_kernel('kernel_Correlation_updateOutput', {
'rbot0': rbot0,
'rbot1': rbot1,
'top': output
}))(
grid=tuple([ output.size(3), output.size(2), output.size(0) ]),
block=tuple([ 32, 1, 1 ]),
shared_mem=first.size(1) * 4,
args=[ n, rbot0.data_ptr(), rbot1.data_ptr(), output.data_ptr() ],
stream=Stream
)
elif first.is_cuda == False:
raise NotImplementedError()
# end
return output
# end上述代码的实现流程如下：
- 创建两个名为 rbot0 和 rbot1 的零张量，用于存储经过重新排列后的输入张量。
- 使用 cupy_launch 启动 CUDA 核函数 kernel_Correlation_rearrange，对输入张量进行重新排列，结果存储在 rbot0 和 rbot1 中。
- 创建名为 output 的零张量，用于存储相关性计算的输出。
- 使用 cupy_launch 启动 CUDA 核函数 kernel_Correlation_updateOutput，计算相关性并将结果存储在 output 中。
- 返回计算得到的输出张量 output。
注意：这段代码中的相关性计算是在 GPU 上进行的，具体实现利用了 CUDA 核函数。代码中还包含了对输入是否在 GPU 上的判断和相应的处理分支。如果输入不在 GPU 上，则抛出 NotImplementedError，表示当前代码不支持在 CPU 上运行。
（7）定义PyTorch 自动求导函数 _FunctionCorrelation 的反向传播逻辑，在 PyTorch 中，通过实现 backward 方法来定义自定义函数的反向传播。
@staticmethod
def backward(self, gradOutput):
first, second, rbot0, rbot1 = self.saved_tensors
assert(gradOutput.is_contiguous() == True)
gradFirst = first.new_zeros([ first.size(0), first.size(1), first.size(2), first.size(3) ]) if self.needs_input_grad[0] == True else None
gradSecond = first.new_zeros([ first.size(0), first.size(1), first.size(2), first.size(3) ]) if self.needs_input_grad[1] == True else None
if first.is_cuda == True:
if gradFirst is not None:
for intSample in range(first.size(0)):
n = first.size(1) * first.size(2) * first.size(3)
cupy_launch('kernel_Correlation_updateGradFirst', cupy_kernel('kernel_Correlation_updateGradFirst', {
'rbot0': rbot0,
'rbot1': rbot1,
'gradOutput': gradOutput,
'gradFirst': gradFirst,
'gradSecond': None
}))(
grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),
block=tuple([ 512, 1, 1 ]),
args=[ n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), gradFirst.data_ptr(), None ],
stream=Stream
)
# end
# end
if gradSecond is not None:
for intSample in range(first.size(0)):
n = first.size(1) * first.size(2) * first.size(3)
cupy_launch('kernel_Correlation_updateGradSecond', cupy_kernel('kernel_Correlation_updateGradSecond', {
'rbot0': rbot0,
'rbot1': rbot1,
'gradOutput': gradOutput,
'gradFirst': None,
'gradSecond': gradSecond
}))(
grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),
block=tuple([ 512, 1, 1 ]),
args=[ n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), None, gradSecond.data_ptr() ],
stream=Stream
)
# end
# end
elif first.is_cuda == False:
raise NotImplementedError()
# end
return gradFirst, gradSecond
# end
# end6.6.4 PWC-Net光流计算
（1）对“pytorch-pwc”存储库中的PWC-Net实现的封装，首先导入了必要的库，如numpy和torch。此外，通过PyTorch设置了一些用于确保可重现性的配置。
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.manual_seed(0)
numpy.random.seed(0)（2）定义函数Backward()对输入张量进行反向运动模糊操作，这个函数的主要目的是通过反向流场对输入张量进行运动模糊。在实现中，首先创建了坐标网格，然后将输入张量与规范化的流场连接起来，最后使用双线性插值对输入张量进行采样。修剪和返回的张量仅包含有效区域，通过应用一个掩码来实现。
def Backward(tensorInput, tensorFlow, device):
# 初始化网格和部分张量
Backward_tensorGrid = {}
Backward_tensorPartial = {}
# 创建水平和垂直方向上的坐标网格
tensorHorizontal = torch.linspace(-1.0, 1.0, tensorFlow.size(3)).view(1, 1, 1, tensorFlow.size(3)).expand(tensorFlow.size(0), -1, tensorFlow.size(2), -1)
tensorVertical = torch.linspace(-1.0, 1.0, tensorFlow.size(2)).view(1, 1, tensorFlow.size(2), 1).expand(tensorFlow.size(0), -1, -1, tensorFlow.size(3))
# 将水平和垂直坐标拼接成一个网格，并存储为Backward_tensorGrid
Backward_tensorGrid[str(tensorFlow.size())] = torch.cat([tensorHorizontal, tensorVertical], 1).to(device)
# 创建部分张量，用于后续拼接
Backward_tensorPartial[str(tensorFlow.size())] = tensorFlow.new_ones([tensorFlow.size(0), 1, tensorFlow.size(2), tensorFlow.size(3)])
# 将流场规范化并与输入张量拼接
tensorFlow = torch.cat([tensorFlow[:, [0], :, :] / ((tensorInput.size(3) - 1.0) / 2.0), tensorFlow[:, [1], :, :] / ((tensorInput.size(2) - 1.0) / 2.0)], 1)
tensorInput = torch.cat([tensorInput, Backward_tensorPartial[str(tensorFlow.size())]], 1)
# 使用双线性插值对输入张量进行采样，并应用流场调整
tensorOutput = torch.nn.functional.grid_sample(input=tensorInput, grid=(Backward_tensorGrid[str(tensorFlow.size())] + tensorFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='zeros')
# 创建一个掩码，用于修剪输出张量
tensorMask = tensorOutput[:, -1:, :, :]
tensorMask[tensorMask > 0.999] = 1.0
tensorMask[tensorMask < 1.0] = 0.0
# 返回修剪后的输出张量
return tensorOutput[:, :-1, :, :] * tensorMask（3）定义特征提取器类Extractor，用于从输入张量中提取多层次的特征。这个类定义了一个卷积神经网络，包含了六个卷积模块（moduleOne到moduleSix），每个模块都由多个卷积层和激活函数（LeakyReLU）组成。在前向传播过程中，输入张量经过每个模块，生成多层次的特征张量，这些特征张量用于后续的光流估计。
class Extractor(torch.nn.Module):
def __init__(self):
super(Extractor, self).__init__()
# 第一层卷积模块
self.moduleOne = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
# 第二层卷积模块
self.moduleTwo = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
# 第三层卷积模块
self.moduleThr = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
# 第四层卷积模块
self.moduleFou = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=64, out_channels=96, kernel_size=3, stride=2, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
# 第五层卷积模块
self.moduleFiv = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=96, out_channels=128, kernel_size=3, stride=2, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
# 第六层卷积模块
self.moduleSix = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=128, out_channels=196, kernel_size=3, stride=2, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
def forward(self, tensorInput):
# 前向传播过程，将输入张量经过每一层卷积模块，得到多层次的特征
tensorOne = self.moduleOne(tensorInput)
tensorTwo = self.moduleTwo(tensorOne)
tensorThr = self.moduleThr(tensorTwo)
tensorFou = self.moduleFou(tensorThr)
tensorFiv = self.moduleFiv(tensorFou)
tensorSix = self.moduleSix(tensorFiv)
# 返回多层次的特征
return [tensorOne, tensorTwo, tensorThr, tensorFou, tensorFiv, tensorSix]（4）定义类Decoder，用于实现光流模型的解码器部分，根据输入特征和前一层的光流信息生成当前层的光流场。解码器负责生成光流场，并且具有多个卷积层来处理输入特征。
class Decoder(torch.nn.Module):
def __init__(self, intLevel):
super(Decoder, self).__init__()
intPrevious = [None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None][intLevel + 1]
intCurrent = [None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None][intLevel + 0]
# 根据光流金字塔的级别选择是否使用反卷积层
if intLevel < 6:
self.moduleUpflow = torch.nn.ConvTranspose2d(in_channels=2, out_channels=2, kernel_size=4, stride=2, padding=1)
if intLevel < 6:
self.moduleUpfeat = torch.nn.ConvTranspose2d(in_channels=intPrevious + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=4, stride=2, padding=1)
if intLevel < 6:
self.dblBackward = [None, None, None, 5.0, 2.5, 1.25, 0.625, None][intLevel + 1]
# 定义卷积模块
self.moduleOne = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=intCurrent, out_channels=128, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
self.moduleTwo = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=intCurrent + 128, out_channels=128, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
self.moduleThr = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=intCurrent + 128 + 128, out_channels=96, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
self.moduleFou = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96, out_channels=64, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
self.moduleFiv = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64, out_channels=32, kernel_size=3, stride=1, padding=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
)
self.moduleSix = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=3, stride=1, padding=1)
)
def forward(self, tensorFirst, tensorSecond, objectPrevious, device):
tensorFlow = None
tensorFeat = None
# 如果是金字塔的第一级，直接进行特征提取
if objectPrevious is None:
tensorFlow = None
tensorFeat = None
tensorVolume = torch.nn.functional.leaky_relu(input=FunctionCorrelation(tensorFirst, tensorSecond, device),
negative_slope=0.1, inplace=False)
tensorFeat = torch.cat([tensorVolume], 1)
# 如果不是金字塔的第一级，则进行特征提取和反向传播
elif objectPrevious is not None:
tensorFlow = self.moduleUpflow(objectPrevious['tensorFlow'])
tensorFeat = self.moduleUpfeat(objectPrevious['tensorFeat'])
tensorVolume = torch.nn.functional.leaky_relu(
input=FunctionCorrelation(tensorFirst=tensorFirst, tensorSecond=Backward(tensorInput=tensorSecond,
tensorFlow=tensorFlow * self.dblBackward,
device=device), device=device),
negative_slope=0.1, inplace=False)
tensorFeat = torch.cat([tensorVolume, tensorFirst, tensorFlow, tensorFeat], 1)
# 进行多层次的特征提取
tensorFeat = torch.cat([self.moduleOne(tensorFeat), tensorFeat], 1)
tensorFeat = torch.cat([self.moduleTwo(tensorFeat), tensorFeat], 1)
tensorFeat = torch.cat([self.moduleThr(tensorFeat), tensorFeat], 1)
tensorFeat = torch.cat([self.moduleFou(tensorFeat), tensorFeat], 1)
tensorFeat = torch.cat([self.moduleFiv(tensorFeat), tensorFeat], 1)
# 通过最后一层卷积得到光流场
tensorFlow = self.moduleSix(tensorFeat)
return {
'tensorFlow': tensorFlow,
'tensorFeat': tensorFeat
}（5）定义类Refiner，用于实现光流模型的细化器部分。细化器负责对之前生成的光流场进行进一步的优化，以提高其准确性。
class Refiner(torch.nn.Module):
def __init__(self):
super(Refiner, self).__init__()
# 定义主要的卷积模块
self.moduleMain = torch.nn.Sequential(
torch.nn.Conv2d(in_channels=81 + 32 + 2 + 2 + 128 + 128 + 96 + 64 + 32, out_channels=128, kernel_size=3, stride=1, padding=1, dilation=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=2, dilation=2),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=4, dilation=4),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=128, out_channels=96, kernel_size=3, stride=1, padding=8, dilation=8),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=96, out_channels=64, kernel_size=3, stride=1, padding=16, dilation=16),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1, dilation=1),
torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
torch.nn.Conv2d(in_channels=32, out_channels=2, kernel_size=3, stride=1, padding=1, dilation=1)
)
def forward(self, tensorInput):
# 通过主要卷积模块得到光流场
return self.moduleMain(tensorInput)（6）定义类PWCNet的，实现了一个光流估计的神经网络，使用了提取器（Extractor）、解码器（Decoder）和细化器（Refiner）等模块。类PWCNet的forward方法接受两个输入图像并返回它们之间的光流场。整个网络包含了提取器、解码器和细化器，以实现精确的光流估计。
class PWCNet(torch.nn.Module):
def __init__(self, pretrain_path):
super(PWCNet, self).__init__()
# 提取器
self.moduleExtractor = Extractor()
# 解码器
self.moduleTwo = Decoder(2)
self.moduleThr = Decoder(3)
self.moduleFou = Decoder(4)
self.moduleFiv = Decoder(5)
self.moduleSix = Decoder(6)
# 细化器
self.moduleRefiner = Refiner()
# 加载预训练模型的权重
self.load_state_dict(torch.load(pretrain_path))
# 设置为评估模式
self.eval()
def forward(self, tensorFirst, tensorSecond, device):
# 对输入图像进行预处理
tensorFirst = tensorFirst[:, [2, 1, 0], :, :] / 255
tensorSecond = tensorSecond[:, [2, 1, 0], :, :] / 255
assert(tensorFirst.size(1) == tensorSecond.size(1))
assert(tensorFirst.size(2) == tensorSecond.size(2))
T, C, intHeight, intWidth = tensorFirst.size()
tensorPreprocessedFirst = tensorFirst.view(T, C, intHeight, intWidth)
tensorPreprocessedSecond = tensorSecond.view(T, C, intHeight, intWidth)
intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))
tensorPreprocessedFirst = torch.nn.functional.interpolate(input=tensorPreprocessedFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
tensorPreprocessedSecond = torch.nn.functional.interpolate(input=tensorPreprocessedSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
tensorPreprocessedFirst = self.moduleExtractor(tensorPreprocessedFirst)
tensorPreprocessedSecond = self.moduleExtractor(tensorPreprocessedSecond)
# 进行光流估计
objectEstimate = self.moduleSix(tensorPreprocessedFirst[-1], tensorPreprocessedSecond[-1], None, device)
objectEstimate = self.moduleFiv(tensorPreprocessedFirst[-2], tensorPreprocessedSecond[-2], objectEstimate, device)
objectEstimate = self.moduleFou(tensorPreprocessedFirst[-3], tensorPreprocessedSecond[-3], objectEstimate, device)
objectEstimate = self.moduleThr(tensorPreprocessedFirst[-4], tensorPreprocessedSecond[-4], objectEstimate, device)
objectEstimate = self.moduleTwo(tensorPreprocessedFirst[-5], tensorPreprocessedSecond[-5], objectEstimate, device)
# 细化光流场
temp = objectEstimate['tensorFlow'] + self.moduleRefiner(objectEstimate['tensorFeat'])
# 进行插值以还原到原始尺寸
tensorFlow = 20.0 * torch.nn.functional.interpolate(input=temp, size=(intHeight, intWidth), mode='bilinear', align_corners=False)
# 将光流场的尺度还原到原始图像的尺度
tensorFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
tensorFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)
return tensorFlow6.6.5 提取I3D特征
（1）定义类ExtractI3D，用于从视频中提取I3D特征。其中forward方法接受视频索引并在每个视频上调用extract方法，该方法处理实际的特征提取工作。提取的特征可以打印出来或保存为NumPy数组，具体取决于on_extraction参数的设置。
class ExtractI3D(torch.nn.Module):
def __init__(self, args):
super(ExtractI3D, self).__init__()
# 从用户输入中构建路径列表等参数
self.path_list = form_list_from_user_input(args)
self.pwc_path = args.pwc_path
self.i3d_rgb_path = args.i3d_rgb_path
self.i3d_flow_path = args.i3d_flow_path
self.min_side_size = args.min_side_size
self.extraction_fps = args.extraction_fps
self.step_size = args.step_size
self.stack_size = args.stack_size
self.show_kinetics_pred = args.show_kinetics_pred
self.kinetics_class_labels = args.kinetics_class_labels
self.keep_frames = args.keep_frames
self.on_extraction = args.on_extraction
self.tmp_path = args.tmp_path
self.output_path = args.output_path
self.progress = tqdm(total=len(self.path_list))
def forward(self, indices: torch.LongTensor):
'''
Arguments:
indices {torch.LongTensor} -- indices to self.path_list
'''
device = indices.device
# 创建PWCNet和I3D模型，并将其移动到指定的设备
pwc_model = PWCNet(self.pwc_path).to(device)
i3d_model = I3D_RGB_FLOW(self.i3d_rgb_path, self.i3d_flow_path).to(device)
for idx in indices:
# 当发生错误时可能会静默地失败，当从Torch数据并行运行时可能会发生
try:
self.extract(device, pwc_model, i3d_model, idx)
except KeyboardInterrupt:
raise KeyboardInterrupt
except Exception as e:
print(e)
print(f'Extraction failed at: {self.path_list[idx]}. Continuing extraction')
# 更新tqdm进度条
self.progress.update()
def extract(self, device: torch.device, pwc_model: torch.nn.Module, i3d_model: torch.nn.Module,
idx: int, video_path: Union[str, None] = None
) -> Dict[str, Union[torch.nn.Module, str]]:
'''The extraction call. Made to clean the forward call a bit.
Arguments:
device {torch.device}
pwc_model {torch.nn.Module}
i3d_model {torch.nn.Module}
idx {int} -- index to self.video_paths
Keyword Arguments:
video_path {Union[str, None]} -- if you would like to use import it and use it as
"path -> i3d features"-fashion (default: {None})
Returns:
Dict[str, Union[torch.nn.Module, str]] -- dict with i3d features and their type
'''
if video_path is None:
video_path = self.path_list[idx]
# 从视频中提取帧，并形成帧的路径列表
frames_dir = extract_frames_from_video(video_path, self.extraction_fps,
self.min_side_size, self.tmp_path)
frame_paths = [os.path.join(frames_dir, fname) for fname in sorted(os.listdir(frames_dir))]
frame_paths = form_iter_list(frame_paths, self.step_size, self.stack_size+1)
# 提取I3D特征
i3d_feats = i3d_features(
frame_paths, self.stack_size, device, pwc_model, i3d_model,
self.show_kinetics_pred, self.kinetics_class_labels
)
# 删除提取帧的文件夹以节省磁盘空间
if not self.keep_frames:
shutil.rmtree(frames_dir)
# 在特征提取完成后执行的操作
if self.on_extraction == 'print':
print(i3d_feats)
elif self.on_extraction == 'save_numpy':
# 如果目录不存在，则创建目录
os.makedirs(self.output_path, exist_ok=True)
# 提取文件名并更改扩展名
filename_rgb = os.path.split(video_path)[-1].replace('.mp4', '_rgb.npy')
filename_flow = os.path.split(video_path)[-1].replace('.mp4', '_flow.npy')
# 构造保存特征的路径
feature_rgb_path = os.path.join(self.output_path, filename_rgb)
feature_flow_path = os.path.join(self.output_path, filename_flow)
# 保存特征
np.save(feature_rgb_path, i3d_feats['rgb'].cpu())
np.save(feature_flow_path, i3d_feats['flow'].cpu())
else:
raise NotImplementedError
return i3d_feats（2）设置一个SimpleNamespace对象，其中包含了参数args，该参数用于初始化ExtractI3D模块。这个模块用于从视频中提取I3D特征。对各个参数的具体说明如下：
- tmp_path：用于存储临时文件的路径。
- keep_frames：一个布尔值，指示是否保留提取的帧文件。
- on_extraction：决定提取完成后的操作，你设置为"save_numpy"，表示将提取的特征保存为NumPy数组。
- pwc_path：PWCNet模型的路径。
- i3d_rgb_path：I3D RGB模型的路径。
- i3d_flow_path：I3D Flow模型的路径。
- min_side_size：视频帧的最小边大小。
- extraction_fps：提取视频帧的帧率。
- stack_size：I3D模型中的帧堆栈大小。
- step_size：提取帧的步幅。
- show_kinetics_pred：是否显示Kinetics预测。
- kinetics_class_labels：Kinetics类别标签的文件路径。
- vggish_model_path：VGGish模型的路径。
- vggish_pca_path：VGGish PCA参数的路径。
- file_with_video_paths：包含视频路径的文件。
- device_ids：用于模型训练的设备ID列表。
- video_paths：视频文件的路径列表。
- output_path：存储提取结果的路径。
（3）在多个设备上并行提取视频特征，并将结果保存到指定的输出路径。
import torch
import argparse
from submodules.video_features.utils.utils import form_list_from_user_input, fix_tensorflow_gpu_allocation
from submodules.video_features.models.vggish.extract_vggish import ExtractVGGish # defined here to avoid import errors
def parallel_feature_extraction(args):
'''Distributes the feature extraction in a embarasingly-parallel fashion. Specifically,
it divides the dataset (list of video paths) among all specified devices evenly.'''
if args.feature_type == 'i3d':
extractor = ExtractI3D(args)
elif args.feature_type == 'vggish':
from models.vggish.extract_vggish import ExtractVGGish # defined here to avoid import errors
fix_tensorflow_gpu_allocation(args)
extractor = ExtractVGGish(args)
video_paths = form_list_from_user_input(args)
indices = torch.arange(len(video_paths))
replicas = torch.nn.parallel.replicate(extractor, args.device_ids[:len(indices)])
inputs = torch.nn.parallel.scatter(indices, args.device_ids[:len(indices)])
torch.nn.parallel.parallel_apply(replicas[:len(inputs)], inputs)
extractor.progress.close()
FEATURES_CACHE_PATH = '/kaggle/working/video_caption_bmt/tmp/'
def extract_video(video_path):
args.feature_type = 'i3d'
args.extraction_fps = 25
args.video_paths = [video_path]
args.output_path = FEATURES_CACHE_PATH
parallel_feature_extraction(args)
args.feature_type = 'vggish'
args.on_extraction = "save_numpy"
args.device_ids = [0]
args.extraction_fps = None
args.video_paths = [video_path]
args.output_path = FEATURES_CACHE_PATH
parallel_feature_extraction(args)对上述代码的具体说明如下：
- ExtractI3D 和 ExtractVGGish 模块：这些模块包含了对于 I3D 和 VGGish 模型的视频特征提取逻辑。这些模块在其他文件中定义。
- parallel_feature_extraction 函数：这个函数用于将视频特征提取任务分配给多个设备进行并行处理。它根据 device_ids 列表将视频路径分配给多个设备，然后在每个设备上运行相应的 ExtractI3D 或 ExtractVGGish 模块。
- extract_video 函数：这个函数是一个简单的包装器，用于提取单个视频的 I3D 和 VGGish 特征。它通过设置一些参数（如视频路径、输出路径等）调用 parallel_feature_extraction 函数。
- FEATURES_CACHE_PATH：存储特征缓存的路径。
（4）下面的代码用于加载模型和其他一些必要的设置，主要是用于加载视频标题生成（caption generation）和视频提议生成（proposal generation）相关的模型。
wk_papth = '/kaggle/working/video_caption_bmt'
prop_generator_model_path = wk_papth + '/sample/best_prop_model.pt'
pretrained_cap_model_path = wk_papth + '/sample/best_cap_model.pt'
max_prop_per_vid = 100
nms_tiou_thresh = 0.4
device_id = 0
cap_cfg, cap_model, train_dataset = load_cap_model(pretrained_cap_model_path, device_id)
prop_cfg, prop_model = load_prop_model(
device_id, prop_generator_model_path, pretrained_cap_model_path, max_prop_per_vid
)（5）对视频实现特征提取工作，具体实现代码如下所示。
video_title = 'file_example_MP4_480_1_5MG'
#%%
video_name = '/kaggle/working/video_caption_bmt/sample/' + video_title
video_path = video_name + '.mp4'
vggish_features_path = FEATURES_CACHE_PATH + video_title + '_vggish.npy'
rgb_features_path = FEATURES_CACHE_PATH + video_title + '_rgb.npy'
flow_features_path = FEATURES_CACHE_PATH + video_title + '_flow.npy'
duration_in_secs = get_video_duration(video_path)
feature_paths = {
'audio': vggish_features_path,
'rgb': rgb_features_path,
'flow': flow_features_path,
}
%cd /kaggle/working/video_caption_bmt/submodules/video_features
extract_video(video_path)
%cd /kaggle/working/video_caption_bmt对上述代码的具体说明如下：
- video_title 和 video_name：分别是视频的标题和完整路径。
- video_path：视频文件的路径，由 video_name 加上文件扩展名 .mp4 构成。
- vggish_features_path、rgb_features_path 和 flow_features_path：分别是 VGGish 音频特征、RGB 图像特征和光流特征的文件路径。
- duration_in_secs：获取视频的时长（以秒为单位）。
- feature_paths：包含了不同类型特征文件路径的字典。
- %cd 命令：用于更改当前工作目录。
- 调用 extract_video 函数：这个函数接受视频路径作为输入，进行特征提取，并将特征保存到指定的路径。在这里，它分别提取了 I3D 特征和 VGGish 特征。
注意：这里的 %cd 命令是用于在 notebook 中切换工作目录的魔术命令。在实际运行时，你可能需要使用相应的命令来切换目录。
（6）如果你使用的是 Jupyter Notebook 或类似的环境，可以直接在代码单元格中执行以下命令查看指定目录中的文件。
!ls /kaggle/working/video_caption_bmt/tmp（7）生成视频的字幕提议，通过先生成提议，然后为每个提议生成字幕描述，最终输出生成的字幕。
# Proposal
proposals = generate_proposals(
prop_model, feature_paths, train_dataset.pad_idx, prop_cfg, device_id, duration_in_secs
)
# NMS if specified
if nms_tiou_thresh is not None:
proposals = non_max_suppresion(proposals.squeeze(), nms_tiou_thresh)
proposals = proposals.unsqueeze(0)
# Captions for each proposal
captions = caption_proposals(
cap_model, feature_paths, train_dataset, cap_cfg, device_id, proposals, duration_in_secs
)
print(captions)上述代码的实现流程如下：
- 使用 prop_model 利用视频的音频、RGB 和光流特征（由 feature_paths 指定）生成初始提议 (proposals)。
- 如果设置了非极大值抑制 (nms_tiou_thresh 不为 None)，则对生成的提议进行非极大值抑制。
- 使用 cap_model 结合音频、RGB 特征、光流特征以及生成的提议，为每个提议生成相应的字幕描述 (captions)。
- 最后，打印输出生成的字幕：
[{'start': 12.3, 'end': 25.4, 'sentence': 'A large bowl of coffee is shown'}, {'start': 0.0, 'end': 4.1, 'sentence': 'A close up of a sink is shown with a title screen'}, {'start': 27.0, 'end': 29.7, 'sentence': 'A close up of a sink is shown'}, {'start': 28.0, 'end': 28.4, 'sentence': 'A close up of a sink is shown'}, {'start': 29.0, 'end': 29.4, 'sentence': 'A close up of a sink is shown'}, {'start': 21.5, 'end': 30.5, 'sentence': 'A large bowl of coffee is shown'}, {'start': 0.3, 'end': 18.2, 'sentence': 'A large bowl of water is shown and leads into a large bowl'}, {'start': 27.1, 'end': 27.5, 'sentence': 'A close up of a sink is shown'}, {'start': 25.1, 'end': 28.7, 'sentence': 'We see a title screen'}, {'start': 28.6, 'end': 30.0, 'sentence': 'A close up of a sink is shown'}]
501
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
