Spaces:
Runtime error
Runtime error
File size: 2,394 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 |
# 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 torch
import torch.nn as nn
from models.utils import LinearELR, Conv2dELR, Downsample2d
class Encoder(torch.nn.Module):
def __init__(self, ninputs, size, nlayers=7, conv=Conv2dELR, lin=LinearELR):
super(Encoder, self).__init__()
self.ninputs = ninputs
height, width = size
self.nlayers = nlayers
ypad = ((height + 2 ** nlayers - 1) // 2 ** nlayers) * 2 ** nlayers - height
xpad = ((width + 2 ** nlayers - 1) // 2 ** nlayers) * 2 ** nlayers - width
self.pad = nn.ZeroPad2d((xpad // 2, xpad - xpad // 2, ypad // 2, ypad - ypad // 2))
self.downwidth = ((width + 2 ** nlayers - 1) // 2 ** nlayers)
self.downheight = ((height + 2 ** nlayers - 1) // 2 ** nlayers)
# compile layers
layers = []
inch, outch = 3, 64
for i in range(nlayers):
layers.append(conv(inch, outch, 4, 2, 1, norm="demod", act=nn.LeakyReLU(0.2)))
if inch == outch:
outch = inch * 2
else:
inch = outch
if outch > 256:
outch = 256
self.down1 = nn.ModuleList([nn.Sequential(*layers)
for i in range(self.ninputs)])
self.down2 = lin(256 * self.ninputs * self.downwidth * self.downheight, 512, norm="demod", act=nn.LeakyReLU(0.2))
self.mu = lin(512, 256)
self.logstd = lin(512, 256)
def forward(self, x, losslist : Optional[List[str]]=None):
assert losslist is not None
x = self.pad(x)
x = [self.down1[i](x[:, i*3:(i+1)*3, :, :]).view(x.size(0), 256 * self.downwidth * self.downheight)
for i in range(self.ninputs)]
x = torch.cat(x, dim=1)
x = self.down2(x)
mu, logstd = self.mu(x) * 0.1, self.logstd(x) * 0.01
if self.training:
z = mu + torch.exp(logstd) * torch.randn(*logstd.size(), device=logstd.device)
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
|