0
0

AI数字人硅基数字人模型训练模型网络结构和训练代码

这种训练的时候加入mask，输出的时候根据mask做处理，直接mask回帖
conv.py

import torch
from torch import nn
from torch.nn import functional as F

class DepthwiseSeparableConv2d(nn.Module):
def init(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
super(DepthwiseSeparableConv2d, self).init()
self.depthwise = nn.Sequential(
nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding,groups=in_channels), # ��Ⱦ���
nn.BatchNorm2d(in_channels),
nn.ReLU6(inplace=True))

  1.     self.pointwise = nn.Sequential(
  2.                         nn.Conv2d(in_channels, out_channels, kernel_size=1),  # ������
  3.                         nn.BatchNorm2d(out_channels),
  4.                         nn.ReLU6(inplace=True))
  6. def forward(self, x):
  7.     x = self.depthwise(x)
  8.     x = self.pointwise(x)
  9.     return x

class Conv2d(nn.Module):
def init(self, cin, cout, kernel_size, stride, padding, residual=False, depth_wise = False, *args, **kwargs):
super().init(*args, **kwargs)

  1.     self.residual = residual
  2.     self.depth_wise =depth_wise
  4.     if depth_wise:
  5.         self.conv_block = DepthwiseSeparableConv2d(cin, cout, kernel_size, stride, padding)
  6.     else:
  7.         self.conv_block = nn.Sequential(
  8.                 nn.Conv2d(cin, cout, kernel_size, stride, padding),
  9.                 nn.BatchNorm2d(cout))
  10.         self.act = nn.ReLU6()
  12. def forward(self, x):
  13.     out = self.conv_block(x)
  14.     if self.residual:
  15.         out += x
  16.     return out if self.depth_wise else self.act(out)

class nonorm_Conv2d(nn.Module):
def init(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs):
super().init(*args, **kwargs)
self.conv_block = nn.Sequential(
nn.Conv2d(cin, cout, kernel_size, stride, padding),
)
self.act = nn.LeakyReLU(0.01, inplace=True)

  1. def forward(self, x):
  2.     out = self.conv_block(x)
  3.     return self.act(out)

class Conv2dTranspose(nn.Module):
def init(self, cin, cout, kernel_size, stride, padding, output_padding=0, *args, **kwargs):
super().init(*args, **kwargs)
self.conv_block = nn.Sequential(
nn.ConvTranspose2d(cin, cout, kernel_size, stride, padding, output_padding),
nn.BatchNorm2d(cout)
)
self.act = nn.ReLU()

  1. def forward(self, x):
  2.     out = self.conv_block(x)
  3.     return self.act(out)

wav2lip.py
import torch
from torch import nn
from torch.nn import functional as F
import math

from models.conv import Conv2dTranspose, Conv2d, nonorm_Conv2d

class Wav2Lip(nn.Module):
def init(self):
super(Wav2Lip, self).init()

  1.     self.face_encoder_blocks = nn.ModuleList([
  2.         nn.Sequential(Conv2d(6, 16, kernel_size=7, stride=1, padding=3, depth_wise=True)), # 96,96       # 288,192    # 144,96
  4.         nn.Sequential(Conv2d(16, 32, kernel_size=3, stride=2, padding=1, depth_wise=True), # 48,48       # 144,96     # 72,48
  5.         Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  6.         Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True)),
  8.         nn.Sequential(Conv2d(32, 64, kernel_size=3, stride=2, padding=1, depth_wise=True),    # 24,24    # 72,48     # 36,24
  9.         Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  10.         Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  11.         Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True)),
  13.         nn.Sequential(Conv2d(64, 128, kernel_size=3, stride=2, padding=1, depth_wise=True),   # 12,12     # 36,24   #18,12
  14.         Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  15.         Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True)),
  17.         nn.Sequential(Conv2d(128, 256, kernel_size=3, stride=2, padding=1, depth_wise=True),       # 6,6   #18,12   #9,6
  18.         Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  19.         Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True)),
  21.         # nn.Sequential(Conv2d(256, 512, kernel_size=3, stride=2, padding=1, depth_wise=True),       # 6,6   #9,6
  22.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  23.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True)),
  25.         # nn.Sequential(Conv2d(512, 512, kernel_size=3, stride=2, padding=1, depth_wise=True),     # 3,3    #5,3
  26.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),),
  27.         #
  28.         # nn.Sequential(Conv2d(512, 512, kernel_size=3, stride=(2, 1), padding=1, depth_wise=True),     # 3,3    #3,3
  29.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),),
  30.         #
  31.         # nn.Sequential(Conv2d(512, 512, kernel_size=3, stride=1, padding=0, depth_wise=True),     # 1, 1
  32.         # Conv2d(512, 512, kernel_size=1, stride=1, padding=0, depth_wise=True)),
  33.         ])
  35.     self.audio_encoder = nn.Sequential(
  36.         Conv2d(1, 32, kernel_size=3, stride=1, padding=1, depth_wise=False),                   #80,16
  37.         Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=False),
  38.         Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=False),
  40.         Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1, depth_wise=False),              #27,16
  41.         Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=False),
  42.         Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=False),
  44.         Conv2d(64, 128, kernel_size=3, stride=3, padding=1, depth_wise=False),                  #9,6
  45.         Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=False),
  46.         Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=False),
  48.         Conv2d(128, 256, kernel_size=3, stride=1, padding=1, depth_wise=False),                 #3,3
  49.         Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=False, depth_wise=False),
  51.         # Conv2d(256, 512, kernel_size=1, stride=1, padding=0, depth_wise=False),                #1,1
  52.         # Conv2d(512, 512, kernel_size=1, stride=1, padding=0, depth_wise=False),                #1,1
  53.         )
  55.     self.face_decoder_blocks = nn.ModuleList([
  56.         nn.Sequential(Conv2d(256, 256, kernel_size=1, stride=1, padding=0, depth_wise=False),),
  58.         # nn.Sequential(Conv2dTranspose(1024, 512, kernel_size=3, stride=1, padding=0), # 3,3
  59.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),),
  60.         #
  61.         # nn.Sequential(Conv2dTranspose(1024, 512, kernel_size=3, stride=(2, 1), padding=1),
  62.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),), # 5, 3
  63.         #
  64.         # nn.Sequential(Conv2dTranspose(1024, 512, kernel_size=3, stride=2, padding=1, output_padding=(0, 1)),
  65.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  66.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),), # 9, 6
  68.         # nn.Sequential(Conv2dTranspose(1024, 512, kernel_size=3, stride=2, padding=1, output_padding=1),
  69.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  70.         # Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),), # 18, 12
  72.         nn.Sequential(Conv2dTranspose(512, 384, kernel_size=3, stride=2, padding=1, output_padding=1),
  73.         Conv2d(384, 384, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  74.         Conv2d(384, 384, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),), # 36, 24
  76.         nn.Sequential(Conv2dTranspose(512, 256, kernel_size=3, stride=2, padding=1, output_padding=1), 
  77.         Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  78.         Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),), # 72, 48
  80.         nn.Sequential(Conv2dTranspose(320, 128, kernel_size=3, stride=2, padding=1, output_padding=1), 
  81.         Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  82.         Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),), # 144, 96
  84.         nn.Sequential(Conv2dTranspose(160, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
  85.         Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),
  86.         Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, depth_wise=True),),]) # 288,192
  88.     self.output_block = nn.Sequential(Conv2d(80, 32, kernel_size=3, stride=1, padding=1, depth_wise=True),
  89.         nn.Conv2d(32, 3, kernel_size=1, stride=1, padding=0),
  90.         nn.Sigmoid()) 
  92. def forward(self, audio_sequences, face_sequences):
  93.     # audio_sequences = (B, T, 1, 80, 16)
  94.     B = audio_sequences.size(0)
  96.     input_dim_size = len(face_sequences.size())
  97.     if input_dim_size > 4:
  98.         audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
  99.         face_sequences = torch.cat([face_sequences[:, :, i] for i in range(face_sequences.size(2))], dim=0)
  101.     audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
  103.     feats = []
  104.     x = face_sequences
  105.     for f in self.face_encoder_blocks:
  106.         x = f(x)
  107.         feats.append(x)
  109.     x = audio_embedding
  110.     for f in self.face_decoder_blocks:
  111.         x = f(x)
  112.         try:
  113.             x = torch.cat((x, feats[-1]), dim=1)
  114.         except Exception as e:
  115.             print(x.size())
  116.             print(feats[-1].size())
  117.             raise e
  118.         
  119.         feats.pop()
  121.     x = self.output_block(x)
  123.     if input_dim_size > 4:
  124.         x = torch.split(x, B, dim=0) # [(B, C, H, W)]
  125.         outputs = torch.stack(x, dim=2) # (B, C, T, H, W)
  127.     else:
  128.         outputs = x
  129.         
  130.     return outputs

class Wav2Lip_disc_qual(nn.Module):
def init(self):
super(Wav2Lip_disc_qual, self).init()

  1.     self.face_encoder_blocks = nn.ModuleList([
  2.         nn.Sequential(nonorm_Conv2d(3, 32, kernel_size=7, stride=1, padding=3)), # 144,192
  4.         nn.Sequential(nonorm_Conv2d(32, 64, kernel_size=5, stride=(1, 2), padding=2), # 144,96
  5.         nonorm_Conv2d(64, 64, kernel_size=5, stride=1, padding=2)),
  7.         nn.Sequential(nonorm_Conv2d(64, 128, kernel_size=5, stride=2, padding=2),    # 72,48
  8.         nonorm_Conv2d(128, 128, kernel_size=5, stride=1, padding=2)),
  10.         nn.Sequential(nonorm_Conv2d(128, 256, kernel_size=5, stride=2, padding=2),   # 36,24
  11.         nonorm_Conv2d(256, 256, kernel_size=5, stride=1, padding=2)),
  13.         nn.Sequential(nonorm_Conv2d(256, 512, kernel_size=3, stride=2, padding=1),       # 18,12
  14.         nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=1)),
  16.         nn.Sequential(nonorm_Conv2d(512, 512, kernel_size=3, stride=2, padding=1),     # 9,6
  17.         nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=1),),
  19.         nn.Sequential(nonorm_Conv2d(512, 512, kernel_size=3, stride=2, padding=1),     # 5,3
  20.         nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=1),),
  22.         nn.Sequential(nonorm_Conv2d(512, 512, kernel_size=3, stride=(2, 1), padding=1),     # 3,3
  23.         nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=1),),
  24.         
  25.         nn.Sequential(nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=0),     # 1, 1
  26.         nonorm_Conv2d(512, 512, kernel_size=1, stride=1, padding=0)),])
  28.     self.binary_pred = nn.Sequential(nn.Conv2d(512, 1, kernel_size=1, stride=1, padding=0), nn.Sigmoid())
  29.     self.label_noise = .0
  31. def get_lower_half(self, face_sequences):
  32.     return face_sequences[:, :, face_sequences.size(2)//2:]
  34. def to_2d(self, face_sequences):
  35.     B = face_sequences.size(0)
  36.     face_sequences = torch.cat([face_sequences[:, :, i] for i in range(face_sequences.size(2))], dim=0)
  37.     return face_sequences
  39. def perceptual_forward(self, false_face_sequences):
  40.     false_face_sequences = self.to_2d(false_face_sequences)
  41.     false_face_sequences = self.get_lower_half(false_face_sequences)
  43.     false_feats = false_face_sequences
  44.     for f in self.face_encoder_blocks:
  45.         false_feats = f(false_feats)
  47.     false_pred_loss = F.binary_cross_entropy(self.binary_pred(false_feats).view(len(false_feats), -1), 
  48.                                     torch.ones((len(false_feats), 1)).cuda())
  50.     return false_pred_loss
  52. def forward(self, face_sequences):
  53.     face_sequences = self.to_2d(face_sequences)
  54.     face_sequences = self.get_lower_half(face_sequences)
  56.     x = face_sequences
  57.     for f in self.face_encoder_blocks:
  58.         x = f(x)
  60.     return self.binary_pred(x).view(len(x), -1)

在这里插入图片描述

请添加图片描述

标签： 人工智能 深度学习 机器学习
本文转载自: https://blog.csdn.net/ebdiy3119/article/details/140275228
版权归原作者 好友hy93119 所有， 如有侵权，请联系我们删除。
登录后发布评论

“AI数字人硅基数字人模型训练模型网络结构和训练代码”的评论:

还没有评论
关于作者
...
overfit同步小助手
文章同步
相关阅读

网络安全法-网络运行安全

使用selenium/drissionpage时如何阻止chrome自动跳转http到https

docker desktop 里部署的Open WebUI 管理员密码忘记了的处理方法

在ubuntu20.04中搭建onsite比赛运行环境

利用开源的低代码表单设计器FcDesigner高效管理和渲染复杂表单结构

Kafka学习笔记

【前端】浏览器输入url到页面呈现发生了什么？

文章导航
没有找到文章内容分段。。。