def __init__(self, zdim, features, channels, lr, betas, eps, device='cpu'):
super().__init__()
self.step = 0
self.alpha = 1.0
self.device = device
self.features = features
self.channels = channels
self.zdim = zdim
self.lr, self.betas, self.eps = lr, betas, eps
in_features = features * 2 ** (8 - max(3, self.step))
out_features = features * 2 ** (8 - max(3, self.step + 1))
self.to_rgb = ll.Conv2d(out_features, self.channels, 1, stride=1, padding=0)
self.prev_to_rgb = None
self.blocks = nn.ModuleList()
self.input_block = nn.Sequential(
ll.ConvBlock(self.zdim, in_features, 4, stride=1, padding=3),
ll.ConvBlock(in_features, in_features, 3, stride=1, padding=1),
ll.ConvBlock(in_features, out_features, 3, stride=1, padding=1),
)
self.optim = optim.Adam(self.parameters(), lr=self.lr, betas=self.betas, eps=self.eps)
self.to(self.device)
def __init__(self, features, channels, lr, betas, eps, device='cpu'):
super().__init__()
self.step = 0
self.alpha = 1.0
self.features = features
self.channels = channels
self.device = device
self.lr, self.betas, self.eps = lr, betas, eps
in_features = features * 2 ** 5
self.from_rgb = ll.Conv2d(channels, in_features, 1, stride=1, padding=0)
self.prev_from_rgb = None
self.blocks = nn.ModuleList()
self.final_block = nn.Sequential(
ll.ConvBlock(in_features, in_features, 3, stride=1, padding=1),
ll.ConvBlock(in_features, in_features, 3, stride=1, padding=1),
ll.MinibatchStddev(),
ll.ConvBlock(in_features+1, in_features, 3, stride=1, padding=1),
ll.ConvBlock(in_features, in_features, 4, stride=1, padding=0),
ll.Conv2d(in_features, 1, 1, stride=1, padding=0),
nn.Flatten(),
)
self.optim = optim.Adam(self.parameters(), lr=self.lr, betas=self.betas, eps=self.eps)
self.to(self.device)
def add_scale(self):
self.step += 1
in_features = self.features * 2 ** (8 - max(3, self.step + 1))
out_features = self.features * 2 ** (8 - max(3, self.step))
self.prev_from_rgb = self.from_rgb
self.from_rgb = ll.Conv2d(self.channels, in_features, 1, stride=1, padding=0)
new_block = nn.Sequential(
ll.ConvBlock(in_features, in_features, 3, stride=1, padding=1),
ll.ConvBlock(in_features, out_features, 3, stride=1, padding=1),
)
self.blocks.insert(0, new_block)
self.optim = optim.Adam(self.parameters(), lr=self.lr, betas=self.betas, eps=self.eps)
self.to(self.device)
def __init__(self):
super(ConvNet, self).__init__()
self.conv_layers = nn.ModuleList()
# Conv layers
# 3 x 84 x 84 ->
self.conv_layers.append(L.ConvBlock(CHANNELS, 32, 8, 4))
self.conv_layers.append(L.ConvBlock(32, 64, 4, 2))
self.conv_layers.append(L.ConvBlock(64, 64, 3, 1))
# -> 64 x 7 x 7
# Flatten and add FC Layers
self.fc_layers = nn.ModuleList()
self.fc_layers.append(nn.Linear(64 * 7 * 7, 256))
self.fc_layers.append(nn.ReLU())
self.fc_layers.append(nn.Linear(256, 448))
self.fc_layers.append(nn.ReLU())