• 视频去噪网络BSVD的实现

    前些天写了视频去噪网络BSVD论文的理解，详情请点击这里，这两个星期动手实践了一下，本篇就来记录一下这个模型的实现。

    这个网络的独特之处在于，它的训练和推理在实现上有所差别。在训练阶段，其使用了TSM（Time Shift Module）结构，而在推理时则使用了BBB（Bidirectional Buffer Block）结构。训练时，网络是一个MIMO（多输入多输出）形式，而在推理时，则将其设计成了单输入、单输出的流式形式。推理时，由于网络中存在16个双向buffer，即BBB，因此，前16帧会输出空数据，16帧之后开始正常输出去噪视频帧，到视频序列结束后，还会继续输出16帧的去噪视频帧，也就是，流式推理整体存在16帧的延迟。这在一些对实时性要求不太高的应用中可以推广，但对于实时性要求严格，并且存储资源有限的应用中，就无法有效应用了。

    下面，我们就通过对官方代码的理解，来聊一聊BSVD的实现。

    官方代码地址：GitHub - ChenyangQiQi/BSVD: [ACM MM 2022] Real-time Streaming Video Denoising with Bidirectional Buffers

    BSVD网络采用了两个UNet级联的方式。

    1. 训练阶段的网络实现

    在训练阶段，网络的实现如下：

    1. class WNet(nn.Module):
    2.     def __init__(self, chns=[32, 64, 128], mid_ch=3, shift_input=False, stage_num=2, in_ch=4, out_ch=3, norm='bn', act='relu', interm_ch=30, blind=False):
    3.     # def __init__(self, chns=[32, 64, 128], mid_ch=3, shift_input=False, stage_num=2, in_ch=4, out_ch=3, norm='bn', act='relu', blind=False):
    4.         super(WNet, self).__init__()
    5.         
    6.         self.stage_num = stage_num
    7.         self.nets_list = nn.ModuleList()
    8.         for i in np.arange(stage_num):
    9.             if i == 0:
    10.                 stage_in_ch = in_ch
    11.             else:
    12.                 stage_in_ch = mid_ch
    13.             if i == (stage_num-1):
    14.                 stage_out_ch = out_ch
    15.             else:
    16.                 stage_out_ch = mid_ch
    17.                 
    18.             # self.nets_list.append(DenBlock(chns=chns, out_ch=stage_out_ch, in_ch=stage_in_ch, shift_input=shift_input, norm=norm, act=act, interm_ch=interm_ch))
    19.             
    20.             if i == 0:
    21.                 self.nets_list.append(DenBlock(chns=chns, out_ch=stage_out_ch, in_ch=stage_in_ch, shift_input=shift_input, norm=norm, act=act, blind=blind, interm_ch=interm_ch))
    22.             else:
    23.                 self.nets_list.append(DenBlock(chns=chns, out_ch=stage_out_ch,
    24.                                            in_ch=stage_in_ch, shift_input=shift_input, norm=norm, act=act, interm_ch=interm_ch))
    25.         # self.temp2 = DenBlock(chns=chns, in_ch=mid_ch, shift_input=shift_input)
    26.  
    27.         # Init weights
    28.         self.reset_params()
    29.  
    30.     @staticmethod
    31.     def weight_init(m):
    32.         if isinstance(m, nn.Conv2d):
    33.             nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
    34.  
    35.     def reset_params(self):
    36.         for _, m in enumerate(self.modules()):
    37.             self.weight_init(m)
    38.  
    39.     def forward(self, x, debug=False):
    40.         # if debug: x_in = x
    41.         # x = self.temp1(x)
    42.         for i in np.arange(self.stage_num):
    43.             if debug: x_temp1 = x
    44.             x = self.nets_list[i](x)
    45.         # if debug: x_temp2 = x
    46.         return x

    网络由两个DenBlock组成，每个DenBlock是一个UNet结构：

    1.  
    2. class DenBlock(nn.Module):
    3.     """ Definition of the denosing block of FastDVDnet.
    4.     Inputs of constructor:
    5.         num_input_frames: int. number of input frames
    6.     Inputs of forward():
    7.         xn: input frames of dim [N, C, H, W], (C=3 RGB)
    8.         noise_map: array with noise map of dim [N, 1, H, W]
    9.     """
    10.  
    11.     def __init__(self, chns=[32, 64, 128], out_ch=3, in_ch=4, shift_input=False, norm='bn', bias=True,  act='relu', interm_ch=30, blind=False):
    12.     # def __init__(self, chns=[32, 64, 128], out_ch=3, in_ch=4, shift_input=False, norm='bn', bias=True,  act='relu', blind=False):
    13.         super(DenBlock, self).__init__()
    14.         self.chs_lyr0, self.chs_lyr1, self.chs_lyr2 = chns
    15.         
    16.         # if stage2: in_ch=3
    17.         if shift_input:
    18.             self.inc = CvBlock(in_ch=in_ch, out_ch=self.chs_lyr0, norm=norm, bias=bias, act=act)
    19.         else:
    20.             self.inc = InputCvBlock(
    21.                 num_in_frames=1, out_ch=self.chs_lyr0, in_ch=in_ch, norm=norm, bias=bias, act=act, interm_ch=interm_ch, blind=blind)
    22.                 # num_in_frames=1, out_ch=self.chs_lyr0, in_ch=in_ch, norm=norm, bias=bias, act=act, blind=blind)
    23.         self.downc0 = DownBlock(in_ch=self.chs_lyr0, out_ch=self.chs_lyr1, norm=norm, bias=bias, act=act)
    24.         self.downc1 = DownBlock(in_ch=self.chs_lyr1, out_ch=self.chs_lyr2, norm=norm, bias=bias, act=act)
    25.         self.upc2 = UpBlock(in_ch=self.chs_lyr2, out_ch=self.chs_lyr1, norm=norm, bias=bias,    act=act)
    26.         self.upc1 = UpBlock(in_ch=self.chs_lyr1, out_ch=self.chs_lyr0, norm=norm, bias=bias,    act=act)
    27.         self.outc = OutputCvBlock(in_ch=self.chs_lyr0, out_ch=out_ch, norm=norm, bias=bias,     act=act)
    28.  
    29.         self.reset_params()
    30.  
    31.     @staticmethod
    32.     def weight_init(m):
    33.         if isinstance(m, nn.Conv2d):
    34.             nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
    35.  
    36.     def reset_params(self):
    37.         for _, m in enumerate(self.modules()):
    38.             self.weight_init(m)
    39.  
    40.     def forward(self, in1):
    41.         '''Args:
    42.             inX: Tensor, [N, C, H, W] in the [0., 1.] range
    43.             noise_map: Tensor [N, 1, H, W] in the [0., 1.] range
    44.         '''
    45.         # Input convolution block
    46.         x0 = self.inc(in1)
    47.         # Downsampling
    48.         x1 = self.downc0(x0)
    49.         x2 = self.downc1(x1)
    50.         # Upsampling
    51.         x2 = self.upc2(x2)
    52.         x1 = self.upc1(x1+x2)
    53.         # Estimation
    54.         x = self.outc(x0+x1)
    55.  
    56.         # Residual
    57.         x[:, :3, :, :] = in1[:, :3, :, :] - x[:, :3, :, :]
    58.  
    59.         return x
    60.

    这段代码与论文中的UNet结构相对应（见下图），包含一个输入层，两个下采样层，两个上采样层，一个输出层。

    输入层没什么特别可说的，主要是两个Conv2d=>BN=>ReLU的组合；输出层也是常规实现，Con2d=>BN=>ReLU=>Con2d，需要注意的是，作者在实现过程中，BN层是没有使用的，是透传通过。

    需要花心思理解的是下采样层和上采样层的实现，因为这两个模块在训练和推理过程中，是有所不同的。

    两个模块的初始实现很简单，定义如下：

    1. class DownBlock(nn.Module):
    2.     '''Downscale + (Conv2d => BN => ReLU)*2'''
    3.  
    4.     def __init__(self, in_ch, out_ch, norm='bn', bias=True, act='relu'):
    5.         super(DownBlock, self).__init__()
    6.         norm_fn = get_norm_function(norm)
    7.         act_fn = get_act_function(act)
    8.         self.convblock = nn.Sequential(
    9.             nn.Conv2d(in_ch, out_ch, kernel_size=3,
    10.                       padding=1, stride=2, bias=bias),
    11.             norm_fn(out_ch),
    12.             act_fn(inplace=True),
    13.             CvBlock(out_ch, out_ch, norm=norm, bias=bias, act=act)
    14.         )
    15.  
    16.     def forward(self, x):
    17.         return self.convblock(x)
    18.  
    19.  
    20. class UpBlock(nn.Module):
    21.     '''(Conv2d => BN => ReLU)*2 + Upscale'''
    22.  
    23.     def __init__(self, in_ch, out_ch, norm='bn', bias=True, act='relu'):
    24.         super(UpBlock, self).__init__()
    25.         # norm_fn = get_norm_function(norm)
    26.         self.convblock = nn.Sequential(
    27.             CvBlock(in_ch, in_ch, norm=norm, bias=bias, act=act),
    28.             nn.Conv2d(in_ch, out_ch*4, kernel_size=3, padding=1, bias=bias),
    29.             nn.PixelShuffle(2)
    30.         )
    31.  
    32.         return self.convblock(x)

    关键在于两者共同调用的子模块CvBlock的实现，在定义时，CvBlock被常规定义为：

    1. class CvBlock(nn.Module):
    2.     '''(Conv2d => BN => ReLU) x 2'''
    3.  
    4.     def __init__(self, in_ch, out_ch, norm='bn', bias=True, act='relu'):
    5.         super(CvBlock, self).__init__()
    6.         norm_fn = get_norm_function(norm)
    7.         act_fn = get_act_function(act)
    8.         self.c1 = nn.Conv2d(in_ch, out_ch, kernel_size=3,
    9.                             padding=1, bias=bias)
    10.         self.b1 = norm_fn(out_ch)
    11.         self.relu1 = act_fn(inplace=True)
    12.         self.c2 = nn.Conv2d(out_ch, out_ch, kernel_size=3,
    13.                             padding=1, bias=bias)
    14.         self.b2 = norm_fn(out_ch)
    15.         self.relu2 = act_fn(inplace=True)
    16.  
    17.     def forward(self, x):
    18.         x = self.c1(x)
    19.         x = self.b1(x)
    20.         x = self.relu1(x)
    21.         x = self.c2(x)
    22.         x = self.b2(x)
    23.         x = self.relu2(x)
    24.         return x

    但接下来，上述定义中的c1和c2则被替换成了TSM实现：

    其中，shift模块的核心实现代码如下，对输入的channels分别向左和向右移动了一定单位（fold）。

    1. def shift(x, n_segment, shift_type, fold_div=3, stride=1, inplace=False):
    2.     nt, c, h, w = x.size()
    3.     n_batch = nt // n_segment
    4.     x = x.view(n_batch, n_segment, c, h, w)
    5.  
    6.     fold = c // fold_div # 32/8 = 4
    7.  
    8.     if inplace:
    9.         # Due to some out of order error when performing parallel computing. 
    10.         # May need to write a CUDA kernel.
    11.         print("WARNING: use inplace shift. it has bugs")
    12.         raise NotImplementedError  
    13.         
    14.     else:
    15.         out = torch.zeros_like(x)
    16.         if not 'toFutureOnly' in shift_type:
    17.             out[:, :-stride, :fold] = x[:, stride:, :fold]  # backward (left shift)
    18.             out[:, stride:, fold: 2 * fold] = x[:, :-stride, fold: 2 * fold]  # forward (right shift)
    19.         else:
    20.             out[:, stride:, : 2 * fold] = x[:, :-stride, : 2 * fold] # right shift only
    21.         out[:, :, 2 * fold:] = x[:, :, 2 * fold:]  # not shift
    22.  
    23.     return out.view(nt, c, h, w)

    2. 推理阶段的网络实现

    在推理阶段，网络实现就显得复杂一些了。大致的网络结构没变，但由于内部的TSM替换成了BBB， 因此没办法严格进行整体网络的加载，只能每一层单独加载训练出来的state_dict。并且，网络推理变成了流式推理，整个网络的定义显得比较凌乱，结构如下：

    1. class BSVD(nn.Module):
    2.     """
    3.         Bidirection-buffer based framework with pipeline-style inference
    4.     """
    5.     def __init__(self, chns=[32, 64, 128], mid_ch=3, shift_input=False, in_ch=4, out_ch=3, norm='bn', act='relu', interm_ch=30, blind=False, 
    6.                  pretrain_ckpt='./experiments/pretrained_ckpt/bsvd-64.pth'):
    7.         super(BSVD, self).__init__()
    8.         self.temp1 = DenBlock(chns=chns, out_ch=mid_ch, in_ch=in_ch,  shift_input=shift_input, norm=norm, act=act, blind=blind, interm_ch=interm_ch)
    9.         self.temp2 = DenBlock(chns=chns, out_ch=out_ch, in_ch=mid_ch, shift_input=shift_input, norm=norm, act=act, blind=blind, interm_ch=interm_ch)
    10.  
    11.         self.shift_num = self.count_shift()
    12.         # Init weights
    13.         self.reset_params()
    14.         if pretrain_ckpt is not None:
    15.             self.load(pretrain_ckpt)
    16.  
    17.     def reset(self):
    18.         self.temp1.reset()
    19.         self.temp2.reset()
    20.     def load(self, path):
    21.         ckpt = torch.load(path)
    22.         print("load from %s"%path)
    23.         ckpt_state = ckpt['params']
    24.         # split the dict here
    25.         if 'module' in list(ckpt_state.keys())[0]:
    26.             base_name = 'module.base_model.'
    27.         else:
    28.             base_name = 'base_model.'
    29.         ckpt_state_1 = extract_dict(ckpt_state, string_name=base_name+'nets_list.0.')
    30.         ckpt_state_2 = extract_dict(ckpt_state, string_name=base_name+'nets_list.1.')
    31.         self.temp1.load_from(ckpt_state_1)
    32.         self.temp2.load_from(ckpt_state_2)
    33.             
    34.     @staticmethod
    35.     def weight_init(m):
    36.         if isinstance(m, nn.Conv2d):
    37.             nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
    38.  
    39.     def reset_params(self):
    40.         for _, m in enumerate(self.modules()):
    41.             self.weight_init(m)
    42.  
    43.     def feedin_one_element(self, x):
    44.         x   = self.temp1(x)
    45.         x   = self.temp2(x)
    46.         return x
    47.     
    48.     def forward(self, input, noise_map=None):
    49.         # N, F, C, H, W -> (N*F, C, H, W)
    50.         if noise_map != None:
    51.             input = torch.cat([input, noise_map], dim=2)
    52.         N, F, C, H, W = input.shape
    53.         input = input.reshape(N*F, C, H, W)
    54.         base_out = self.streaming_forward(input)
    55.         NF, C, H, W = base_out.shape
    56.         base_out = base_out.reshape(N, F, C, H, W)
    57.         return base_out
    58.     
    59.     def streaming_forward(self, input_seq):
    60.         """
    61.         pipeline-style inference
    63.         Args:
    64.             Noisy video stream
    66.         Returns:
    67.             Denoised video stream
    68.         """
    69.         out_seq = []
    70.         if isinstance(input_seq, torch.Tensor):
    71.             n,c,h,w = input_seq.shape
    72.             input_seq = [input_seq[i:i+1, ...] for i in np.arange(n)]
    73.         assert type(input_seq) == list, "convert the input into a sequence"
    74.         _,c,h,w = input_seq[0].shape
    75.         with torch.no_grad():
    76.             for i, x in enumerate(input_seq):
    77.  
    78.                 x_cuda = x.cuda()
    79.                 x_cuda = self.feedin_one_element(x_cuda)
    80.                 # if x_cuda is not None: x_cuda = x_cuda.cpu()
    81.                 if isinstance(x_cuda, torch.Tensor):
    82.                     out_seq.append(x_cuda)
    83.                 else:
    84.                     out_seq.append(x_cuda)
    85.  
    86.             end_out = self.feedin_one_element(None)
    87.  
    88.             out_seq.append(end_out)
    89.  
    90.             # end stage
    91.             while 1:
    92.                 end_out = self.feedin_one_element(None)
    93.                 
    94.                 if len(out_seq) == (self.shift_num+len(input_seq)):
    95.                     break
    96.  
    97.                 out_seq.append(end_out)
    98.  
    99.             # number of temporal shift is 2, last element is 0
    100.             # TODO fix init and end frames
    101.             out_seq_clip = out_seq[self.shift_num:]
    102.             self.reset()
    103.             return torch.cat(out_seq_clip, dim=0)
    104.  
    105.     def count_shift(self):
    106.         count = 0
    107.         for name, module in self.named_modules():
    108.             # print(type(module))
    109.             if "BiBufferConv" in str(type(module)):
    110.                 count+=1
    111.         return count

    两个UNet的定义（DenBlock）大体上没发生变化，但下采样模块和上采样模块的定义发生了改变。

    下采样层如下，原来带有TSM的CvBlock换成了MemCvBlock：

    上采样模块也类似：

     

    而MemCvBlock则调用了BBB模块，BBB模块的实现如下，这是整个算法的核心：

    1. class BiBufferConv(nn.Module):
    2.     def __init__(self,
    3.             in_channels,
    4.             out_channels,
    5.             kernel_size,
    6.             stride=1,
    7.             padding=0,
    8.             bias=True
    9.         ) -> None:
    10.         super(BiBufferConv, self).__init__()
    11.         self.op = ShiftConv(
    12.             in_channels,
    13.             out_channels,
    14.             kernel_size,
    15.             stride,
    16.             padding,
    17.             bias
    18.         )
    19.         self.out_channels = out_channels
    20.         self.left_fold_2fold = None
    21.         # self.zero_tensor = None
    22.         self.center = None
    23.         
    24.     def reset(self):
    25.         self.left_fold_2fold = None
    26.         self.center = None
    27.         
    28.     def forward(self, input_right, verbose=False):
    29.         fold_div = 8
    30.         if input_right is not None:
    31.             self.n, self.c, self.h, self.w = input_right.size()
    32.             self.fold = self.c//fold_div
    33.         # Case1: In the start or end stage, the memory is empty
    34.         if self.center is None:
    35.             self.center = input_right
    36.             # if verbose:
    37.             
    38.             if input_right is not None:
    39.                 if self.left_fold_2fold is None:
    40.                     # In the start stage, the memory and left tensor is empty
    41.  
    42.                     self.left_fold_2fold = torch.zeros((self.n, self.fold, self.h, self.w), device=torch.device('cuda'))
    43.                 if verbose: print("%f+none+%f = none"%(torch.mean(self.left_fold_2fold), torch.mean(input_right)))
    44.             else:
    45.                 # in the end stage, both feed in and memory are empty
    46.                 if verbose: print("%f+none+none = none"%(torch.mean(self.left_fold_2fold)))
    47.                 # print("self.center is None")
    48.             return None
    49.         # Case2: Center is not None, but input_right is None
    50.         elif input_right is None:
    51.             # In the last procesing stage, center is 0
    52.             output =  self.op(self.left_fold_2fold, self.center, torch.zeros((self.n, self.fold, self.h, self.w), device=torch.device('cuda')))
    53.             if verbose: print("%f+%f+none = %f"%(torch.mean(self.left_fold_2fold), torch.mean(self.center), torch.mean(output)))
    54.         else:
    55.             
    56.             output =  self.op(self.left_fold_2fold, self.center, input_right)
    57.             if verbose: print("%f+%f+%f = %f"%(torch.mean(self.left_fold_2fold), torch.mean(self.center), torch.mean(input_right), torch.mean(output)))
    58.             # if output == 57:
    59.                 # a = 1
    60.         self.left_fold_2fold = self.center[:, self.fold:2*self.fold, :, :]
    61.         self.center = input_right
    62.         return output

    这样，通过BBB模块，就实现了16个双向Buffer的填充、更新和清空。

    限于篇幅，先梳理出个大体的思路，实际上还有很多细节需要特别关注，留待下一篇来写吧。

  • 相关阅读:
    企业级Java EE架构设计精深实践
    智慧安防/视频分析云平台EasyCVR不显示告警图片该如何解决？
    springBoot依赖管理机制
    unity UnityWebRequest application/x-www-form-urlencoded上传字符对
    手部数据太难找？最全手部开源数据集分享
    基于Java+SpringBoot+Vue前后端分离医疗报销系统设计和实现
    命令模式（Command Pattern）
    Cocos Creator3.8 项目实战（九）2D UI DrawCall优化详解（下）
    【vscode编辑器插件】前端 php unity自用插件分享
    医院预约挂号系统,java医院预约挂号系统,医院预约挂号管理系统毕业设计作品
  • 原文地址：https://blog.csdn.net/DeliaPu/article/details/133997692