def on_build(self, in_ch, ae_ch, ae_out_ch, **kwargs):
self.ae_out_ch = ae_out_ch
self.dense_norm = nn.DenseNorm()
self.dense1 = nn.Dense(in_ch, ae_ch)
self.dense2 = nn.Dense(
ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch)
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def on_build(self):
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
self.dense1 = nn.Dense(in_ch, ae_ch)
self.dense2 = nn.Dense(
ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch)
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def on_build(self):
in_ch, lowest_dense_res, ae_ch, ae_out_ch, d_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch, self.d_ch
self.dense1 = nn.Dense( in_ch, ae_ch, kernel_initializer=tf.initializers.orthogonal )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch, maxout_features=4, kernel_initializer=tf.initializers.orthogonal )
self.upscale1 = Upscale(ae_out_ch, d_ch*8)
self.res1 = ResidualBlock(d_ch*8)
def on_build(self):
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
if 'u' in opts:
self.dense_norm = nn.DenseNorm()
self.dense1 = nn.Dense( in_ch, ae_ch )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def on_build(self):
self.conv1 = nn.Conv2D (3, 64, kernel_size=3, strides=1, padding='SAME')
self.dense1 = nn.Dense (1, 64, use_bias=False)
self.dense2 = nn.Dense (1, 64, use_bias=False)
self.e0_conv0 = nn.Conv2D (64, 64, kernel_size=3, strides=1, padding='SAME')
self.e0_conv1 = nn.Conv2D (64, 64, kernel_size=3, strides=1, padding='SAME')
self.e1_conv0 = nn.Conv2D (64, 112, kernel_size=3, strides=1, padding='SAME')
self.e1_conv1 = nn.Conv2D (112, 112, kernel_size=3, strides=1, padding='SAME')
self.e2_conv0 = nn.Conv2D (112, 192, kernel_size=3, strides=1, padding='SAME')
self.e2_conv1 = nn.Conv2D (192, 192, kernel_size=3, strides=1, padding='SAME')
self.e3_conv0 = nn.Conv2D (192, 336, kernel_size=3, strides=1, padding='SAME')
self.e3_conv1 = nn.Conv2D (336, 336, kernel_size=3, strides=1, padding='SAME')
self.e4_conv0 = nn.Conv2D (336, 512, kernel_size=3, strides=1, padding='SAME')
self.e4_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
self.center_conv0 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
self.center_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
self.center_conv2 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
self.center_conv3 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
self.d4_conv0 = nn.Conv2D (1024, 512, kernel_size=3, strides=1, padding='SAME')
self.d4_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
self.d3_conv0 = nn.Conv2D (848, 512, kernel_size=3, strides=1, padding='SAME')
self.d3_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
self.d2_conv0 = nn.Conv2D (704, 288, kernel_size=3, strides=1, padding='SAME')
self.d2_conv1 = nn.Conv2D (288, 288, kernel_size=3, strides=1, padding='SAME')
self.d1_conv0 = nn.Conv2D (400, 160, kernel_size=3, strides=1, padding='SAME')
self.d1_conv1 = nn.Conv2D (160, 160, kernel_size=3, strides=1, padding='SAME')
self.d0_conv0 = nn.Conv2D (224, 96, kernel_size=3, strides=1, padding='SAME')
self.d0_conv1 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')
self.out1x_conv0 = nn.Conv2D (96, 48, kernel_size=3, strides=1, padding='SAME')
self.out1x_conv1 = nn.Conv2D (48, 3, kernel_size=3, strides=1, padding='SAME')
self.dec2x_conv0 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')
self.dec2x_conv1 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')
self.out2x_conv0 = nn.Conv2D (96, 48, kernel_size=3, strides=1, padding='SAME')
self.out2x_conv1 = nn.Conv2D (48, 3, kernel_size=3, strides=1, padding='SAME')
self.dec4x_conv0 = nn.Conv2D (96, 72, kernel_size=3, strides=1, padding='SAME')
self.dec4x_conv1 = nn.Conv2D (72, 72, kernel_size=3, strides=1, padding='SAME')
self.out4x_conv0 = nn.Conv2D (72, 36, kernel_size=3, strides=1, padding='SAME')
self.out4x_conv1 = nn.Conv2D (36, 3 , kernel_size=3, strides=1, padding='SAME')
def on_build(self, e_ch, levels):
self.enc_blocks = {}
self.from_rgbs = {}
self.dense_norm = nn.DenseNorm()
for level in range(levels, -1, -1):
self.from_rgbs[level] = FromRGB(level_chs[level])
if level != 0:
self.enc_blocks[level] = EncBlock(
level_chs[level], level_chs[level - 1], level)
self.ae_dense1 = nn.Dense(ae_res * ae_res * ae_ch, 256)
self.ae_dense2 = nn.Dense(256, ae_res * ae_res * ae_ch)
def on_build(self):
self.down1 = Downscale(input_ch, e_dims, kernel_size=5)
self.res1 = ResidualBlock(e_dims)
self.down2 = Downscale(e_dims, e_dims*2, kernel_size=5)
self.down3 = Downscale(e_dims*2, e_dims*4, kernel_size=5)
self.down4 = Downscale(e_dims*4, e_dims*8, kernel_size=5)
self.down5 = Downscale(e_dims*8, e_dims*8, kernel_size=5)
self.res5 = ResidualBlock(e_dims*8)
self.dense1 = nn.Dense( (( resolution//(2**5) )**2) * e_dims*8, ae_dims )
def on_build(self):
global lowest_dense_res
in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch
if 'u' in opts:
self.dense_norm = nn.DenseNorm()
self.dense = []
last_ch = in_ch
for i in range(self.num_layers):
self.dense.append(nn.Dense(last_ch, ae_ch))
last_ch = ae_ch
self.dense.append(
nn.Dense(
ae_ch,
lowest_dense_res * lowest_dense_res * ae_out_ch))
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def on_build(self, in_ch, e_ch, ae_ch):
self.down1 = Downscale(in_ch, e_ch, kernel_size=5)
self.res1 = ResidualBlock(e_ch)
self.down2 = Downscale(e_ch, e_ch*2, kernel_size=5)
self.down3 = Downscale(e_ch*2, e_ch*4, kernel_size=5)
self.down4 = Downscale(e_ch*4, e_ch*8, kernel_size=5)
self.down5 = Downscale(e_ch*8, e_ch*8, kernel_size=5)
self.res5 = ResidualBlock(e_ch*8)
self.dense1 = nn.Dense( lowest_dense_res*lowest_dense_res*e_ch*8, ae_ch )
def on_build(self, in_ch, ae_ch):
self.dense1 = nn.Dense(in_ch, ae_ch)
def on_build(self, ae_ch):
self.dense1 = nn.Dense( ae_ch+1, 1024 )
self.dense2 = nn.Dense( 1024, 2048 )
self.dense3 = nn.Dense( 2048, 4096 )
self.dense4 = nn.Dense( 4096, 4096 )
self.dense5 = nn.Dense( 4096, ae_ch )
def on_build(self, in_ch, base_ch, out_ch):
class ConvBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch):
self.conv = nn.Conv2D(in_ch,
out_ch,
kernel_size=3,
padding='SAME')
self.frn = nn.FRNorm2D(out_ch)
self.tlu = nn.TLU(out_ch)
def forward(self, x):
x = self.conv(x)
x = self.frn(x)
x = self.tlu(x)
return x
class UpConvBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch):
self.conv = nn.Conv2DTranspose(in_ch,
out_ch,
kernel_size=3,
padding='SAME')
self.frn = nn.FRNorm2D(out_ch)
self.tlu = nn.TLU(out_ch)
def forward(self, x):
x = self.conv(x)
x = self.frn(x)
x = self.tlu(x)
return x
self.base_ch = base_ch
self.conv01 = ConvBlock(in_ch, base_ch)
self.conv02 = ConvBlock(base_ch, base_ch)
self.bp0 = nn.BlurPool(filt_size=4)
self.conv11 = ConvBlock(base_ch, base_ch * 2)
self.conv12 = ConvBlock(base_ch * 2, base_ch * 2)
self.bp1 = nn.BlurPool(filt_size=3)
self.conv21 = ConvBlock(base_ch * 2, base_ch * 4)
self.conv22 = ConvBlock(base_ch * 4, base_ch * 4)
self.bp2 = nn.BlurPool(filt_size=2)
self.conv31 = ConvBlock(base_ch * 4, base_ch * 8)
self.conv32 = ConvBlock(base_ch * 8, base_ch * 8)
self.conv33 = ConvBlock(base_ch * 8, base_ch * 8)
self.bp3 = nn.BlurPool(filt_size=2)
self.conv41 = ConvBlock(base_ch * 8, base_ch * 8)
self.conv42 = ConvBlock(base_ch * 8, base_ch * 8)
self.conv43 = ConvBlock(base_ch * 8, base_ch * 8)
self.bp4 = nn.BlurPool(filt_size=2)
self.conv51 = ConvBlock(base_ch * 8, base_ch * 8)
self.conv52 = ConvBlock(base_ch * 8, base_ch * 8)
self.conv53 = ConvBlock(base_ch * 8, base_ch * 8)
self.bp5 = nn.BlurPool(filt_size=2)
self.dense1 = nn.Dense(4 * 4 * base_ch * 8, 512)
self.dense2 = nn.Dense(512, 4 * 4 * base_ch * 8)
self.up5 = UpConvBlock(base_ch * 8, base_ch * 4)
self.uconv53 = ConvBlock(base_ch * 12, base_ch * 8)
self.uconv52 = ConvBlock(base_ch * 8, base_ch * 8)
self.uconv51 = ConvBlock(base_ch * 8, base_ch * 8)
self.up4 = UpConvBlock(base_ch * 8, base_ch * 4)
self.uconv43 = ConvBlock(base_ch * 12, base_ch * 8)
self.uconv42 = ConvBlock(base_ch * 8, base_ch * 8)
self.uconv41 = ConvBlock(base_ch * 8, base_ch * 8)
self.up3 = UpConvBlock(base_ch * 8, base_ch * 4)
self.uconv33 = ConvBlock(base_ch * 12, base_ch * 8)
self.uconv32 = ConvBlock(base_ch * 8, base_ch * 8)
self.uconv31 = ConvBlock(base_ch * 8, base_ch * 8)
self.up2 = UpConvBlock(base_ch * 8, base_ch * 4)
self.uconv22 = ConvBlock(base_ch * 8, base_ch * 4)
self.uconv21 = ConvBlock(base_ch * 4, base_ch * 4)
self.up1 = UpConvBlock(base_ch * 4, base_ch * 2)
self.uconv12 = ConvBlock(base_ch * 4, base_ch * 2)
self.uconv11 = ConvBlock(base_ch * 2, base_ch * 2)
self.up0 = UpConvBlock(base_ch * 2, base_ch)
self.uconv02 = ConvBlock(base_ch * 2, base_ch)
self.uconv01 = ConvBlock(base_ch, base_ch)
self.out_conv = nn.Conv2D(base_ch,
out_ch,
kernel_size=3,
padding='SAME')
def on_build(self):
self.dense2 = nn.Dense(ae_dims, inter_res * inter_res * inter_dims)
def on_build(self, in_ch, out_ch, level):
self.zero_level = level == 0
self.dense1 = nn.Dense()
def on_build(self):
in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
self.dense1 = nn.Dense( in_ch, ae_ch, kernel_initializer=tf.initializers.orthogonal )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch, kernel_initializer=tf.initializers.orthogonal )
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def on_build(self):
ae_ch, ae_out_ch = self.ae_ch, self.ae_out_ch
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
def on_build(self, in_ch, n_classes):
self.dense1 = nn.Dense(in_ch, 4096)
self.dense2 = nn.Dense(4096, 4096)
self.pitch_dense = nn.Dense(4096, n_classes)
self.yaw_dense = nn.Dense(4096, n_classes)