-
Notifications
You must be signed in to change notification settings - Fork 0
/
decoder.py
37 lines (30 loc) · 1.28 KB
/
decoder.py
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
import torch.nn as nn
from helper import ResidualBlock, NonLocalBlock, UpSampleBlock, GroupNorm, Swish
class Decoder(nn.Module):
def __init__(self, args):
super(Decoder, self).__init__()
channels = [512, 256, 256, 128, 128]
attn_resolutions = [16]
num_res_blocks = 3
resolution = 16
in_channels = channels[0]
layers = [nn.Conv2d(args.latent_dim, in_channels, 3, 1, 1),
ResidualBlock(in_channels, in_channels),
NonLocalBlock(in_channels),
ResidualBlock(in_channels, in_channels)]
for i in range(len(channels)):
out_channels = channels[i]
for j in range(num_res_blocks):
layers.append(ResidualBlock(in_channels, out_channels))
in_channels = out_channels
if resolution in attn_resolutions:
layers.append(NonLocalBlock(in_channels))
if i != 0:
layers.append(UpSampleBlock(in_channels))
resolution *= 2
layers.append(GroupNorm(in_channels))
layers.append(Swish())
layers.append(nn.Conv2d(in_channels, args.image_channels, 3, 1, 1))
self.model = nn.Sequential(*layers)
def forward(self, x):
return self.model(x)