|
import torch |
|
from torch import nn |
|
from torch.nn import functional as F |
|
|
|
from .conv import Conv2d |
|
|
|
class SyncNet_color(nn.Module): |
|
def __init__(self): |
|
super(SyncNet_color, self).__init__() |
|
|
|
self.face_encoder = nn.Sequential( |
|
Conv2d(15, 32, kernel_size=(7, 7), stride=1, padding=3), |
|
|
|
Conv2d(32, 64, kernel_size=5, stride=(1, 2), padding=1), |
|
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), |
|
|
|
Conv2d(64, 128, kernel_size=3, stride=2, padding=1), |
|
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), |
|
|
|
Conv2d(128, 256, kernel_size=3, stride=2, padding=1), |
|
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), |
|
|
|
Conv2d(256, 512, kernel_size=3, stride=2, padding=1), |
|
Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True), |
|
|
|
Conv2d(512, 512, kernel_size=3, stride=2, padding=1), |
|
Conv2d(512, 512, kernel_size=3, stride=1, padding=0), |
|
Conv2d(512, 512, kernel_size=1, stride=1, padding=0),) |
|
|
|
self.audio_encoder = nn.Sequential( |
|
Conv2d(1, 32, kernel_size=3, stride=1, padding=1), |
|
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), |
|
|
|
Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1), |
|
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), |
|
|
|
Conv2d(64, 128, kernel_size=3, stride=3, padding=1), |
|
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), |
|
|
|
Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1), |
|
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), |
|
Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), |
|
|
|
Conv2d(256, 512, kernel_size=3, stride=1, padding=0), |
|
Conv2d(512, 512, kernel_size=1, stride=1, padding=0),) |
|
|
|
def forward(self, audio_sequences, face_sequences): |
|
face_embedding = self.face_encoder(face_sequences) |
|
audio_embedding = self.audio_encoder(audio_sequences) |
|
|
|
audio_embedding = audio_embedding.view(audio_embedding.size(0), -1) |
|
face_embedding = face_embedding.view(face_embedding.size(0), -1) |
|
|
|
audio_embedding = F.normalize(audio_embedding, p=2, dim=1) |
|
face_embedding = F.normalize(face_embedding, p=2, dim=1) |
|
|
|
|
|
return audio_embedding, face_embedding |
|
|
