deepkyu's picture
initial commit
1ba3df3
raw
history blame
2.46 kB
import torch
import torch.nn as nn
from models.module import Conv2d, StyleAttentionBlock
_ENCODER_CHANNEL_DEFAULT = 256
class Encoder(nn.Module):
def __init__(self, hp, in_channels=1, out_channels=_ENCODER_CHANNEL_DEFAULT):
super().__init__()
self.hp = hp
self.module = nn.ModuleList()
def forward(self, x):
for block in self.module:
x = block(x)
return x
class ContentVanillaEncoder(Encoder):
def __init__(self, hp, in_channels, out_channels):
super().__init__(hp, in_channels, out_channels)
self.depth = hp.encoder.content.depth
assert out_channels // (2 ** self.depth) >= in_channels * 2, "Output channel should be increased"
self.module = nn.ModuleList()
self.module.append(
Conv2d(in_channels, out_channels // (2 ** self.depth),
kernel_size=7, padding=3, padding_mode='reflect', bias=False)
)
for layer_idx in range(1, self.depth + 1): # downsample
self.module.append(
Conv2d(out_channels // (2 ** (self.depth - layer_idx + 1)),
out_channels // (2 ** (self.depth - layer_idx)),
kernel_size=3, stride=2, padding=1, bias=False)
)
class StyleVanillaEncoder(Encoder):
def __init__(self, hp, in_channels, out_channels):
super().__init__(hp, in_channels, out_channels)
self.depth = hp.encoder.style.depth
assert out_channels // (2 ** self.depth) >= in_channels * 2, "Output channel should be increased"
encoder_module = []
encoder_module.append(
Conv2d(in_channels, out_channels // (2 ** self.depth),
kernel_size=7, padding=3, padding_mode='reflect', bias=False)
)
for layer_idx in range(1, self.depth + 1): # downsample
encoder_module.append(
Conv2d(out_channels // (2 ** (self.depth - layer_idx + 1)),
out_channels // (2 ** (self.depth - layer_idx)),
kernel_size=3, stride=2, padding=1, bias=False)
)
self.add_module("encoder_module", nn.Sequential(*encoder_module))
self.add_module("attention_module", StyleAttentionBlock(out_channels))
def forward(self, x):
B, K, H, W = x.size()
out = self.encoder_module(x.view(-1, 1, H, W))
out = self.attention_module(out, B, K)
return out