File size: 2,382 Bytes
81ecb2b |
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 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from typing import Optional, List
import numpy as np
import torch
import torch.nn as nn
from models.utils import LinearELR, Conv2dELR
class Encoder(torch.nn.Module):
def __init__(self, latentdim=256, hiq=True, texin=True,
conv=Conv2dELR, lin=LinearELR,
demod=True, texsize=1024, vertsize=21918):
super(Encoder, self).__init__()
self.latentdim = latentdim
self.vertbranch = lin(vertsize, 256, norm="demod", act=nn.LeakyReLU(0.2))
if texin:
cm = 2 if hiq else 1
layers = []
chout = 128*cm
chin = 128*cm
nlayers = int(np.log2(texsize)) - 2
for i in range(nlayers):
if i == nlayers - 1:
chin = 3
layers.append(
conv(chin, chout, 4, 2, 1, norm="demod" if demod else None, act=nn.LeakyReLU(0.2)))
if chin == chout:
chin = chout // 2
else:
chout = chin
self.texbranch1 = nn.Sequential(*(layers[::-1]))
self.texbranch2 = lin(cm*128*4*4, 256, norm="demod", act=nn.LeakyReLU(0.2))
self.mu = lin(512, self.latentdim)
self.logstd = lin(512, self.latentdim)
else:
self.mu = lin(256, self.latentdim)
self.logstd = lin(256, self.latentdim)
def forward(self, verts, texture : Optional[torch.Tensor]=None, losslist : Optional[List[str]]=None):
assert losslist is not None
x = self.vertbranch(verts.view(verts.size(0), -1))
if texture is not None:
texture = self.texbranch1(texture).reshape(verts.size(0), -1)
texture = self.texbranch2(texture)
x = torch.cat([x, texture], dim=1)
mu, logstd = self.mu(x) * 0.1, self.logstd(x) * 0.01
if self.training:
z = mu + torch.exp(logstd) * torch.randn_like(logstd)
else:
z = mu
losses = {}
if "kldiv" in losslist:
losses["kldiv"] = torch.mean(-0.5 - logstd + 0.5 * mu ** 2 + 0.5 * torch.exp(2 * logstd), dim=-1)
return {"encoding": z}, losses
|