def __init__(self, config):
super(NetworkCrop, self).__init__()
latent_size = config['latent_obj_size']
self.register_buffer('prior_mu', torch.zeros([latent_size]))
self.register_buffer('prior_logvar', torch.zeros([latent_size]))
self.enc_pres = LinearBlock(
hidden_list=config['enc_pres_hidden'],
in_features=config['state_full_size'] + config['state_crop_size'],
out_features=3,
)
self.enc_obj = LinearBlock(
hidden_list=config['enc_obj_hidden'],
in_features=config['state_crop_size'],
out_features=latent_size * 2,
)
self.dec_apc = DecoderResidual(
avg_hidden_list_rev=config['dec_apc_avg_hidden_rev'],
res_channel_list_rev=config['dec_apc_res_channel_rev'],
res_kernel_list_rev=config['dec_apc_res_kernel_rev'],
res_stride_list_rev=config['dec_apc_res_stride_rev'],
res_hidden_list_rev=config['dec_apc_res_hidden_rev'],
in_features=latent_size,
image_shape=config['crop_shape'],
)
self.dec_shp = DecoderBlock(
channel_list_rev=config['dec_shp_channel_rev'],
kernel_list_rev=config['dec_shp_kernel_rev'],
stride_list_rev=config['dec_shp_stride_rev'],
hidden_list_rev=config['dec_shp_hidden_rev'],
in_features=latent_size,
out_shape=[1, *config['crop_shape'][1:]],
)
def __init__(self, config):
super(DecoderObj, self).__init__()
self.dec_apc = DecoderBlock(
channel_list_rev=config['dec_apc_channel_rev'],
kernel_list_rev=config['dec_apc_kernel_rev'],
stride_list_rev=config['dec_apc_stride_rev'],
hidden_list_rev=config['dec_apc_hidden_rev'],
in_features=config['state_obj_size'],
out_shape=config['image_shape'],
)
self.dec_shp = DecoderBlock(
channel_list_rev=config['dec_shp_channel_rev'],
kernel_list_rev=config['dec_shp_kernel_rev'],
stride_list_rev=config['dec_shp_stride_rev'],
hidden_list_rev=config['dec_shp_hidden_rev'],
in_features=config['state_obj_size'],
out_shape=[1, *config['image_shape'][1:]],
)
def __init__(self, config):
super(DecoderBck, self).__init__()
self.dec = DecoderBlock(
channel_list_rev=config['dec_bck_channel_rev'],
kernel_list_rev=config['dec_bck_kernel_rev'],
stride_list_rev=config['dec_bck_stride_rev'],
hidden_list_rev=config['dec_bck_hidden_rev'],
in_features=config['state_bck_size'],
out_shape=config['image_shape'],
)
def __init__(self, avg_hidden_list_rev, res_channel_list_rev,
res_kernel_list_rev, res_stride_list_rev, res_hidden_list_rev,
in_features, image_shape):
super(DecoderResidual, self).__init__()
self.dec_avg = LinearBlock(
hidden_list=reversed(avg_hidden_list_rev),
in_features=in_features,
out_features=image_shape[0],
)
self.dec_res = DecoderBlock(
channel_list_rev=res_channel_list_rev,
kernel_list_rev=res_kernel_list_rev,
stride_list_rev=res_stride_list_rev,
hidden_list_rev=res_hidden_list_rev,
in_features=in_features,
out_shape=image_shape,
)
def __init__(self, config):
super(NetworkWhat, self).__init__()
latent_size = config['latent_what_size']
self.register_buffer('prior_mu', torch.zeros([latent_size]))
self.register_buffer('prior_logvar', torch.zeros([latent_size]))
self.enc = EncoderBlock(
channel_list=config['enc_what_channel'],
kernel_list=config['enc_what_kernel'],
stride_list=config['enc_what_stride'],
hidden_list=config['enc_what_hidden'],
in_shape=[config['crop_shape'][0] * 2, *config['crop_shape'][1:]],
out_features=latent_size * 2,
)
self.dec = DecoderBlock(
channel_list_rev=config['dec_what_channel_rev'],
kernel_list_rev=config['dec_what_kernel_rev'],
stride_list_rev=config['dec_what_stride_rev'],
hidden_list_rev=config['dec_what_hidden_rev'],
in_features=latent_size,
out_shape=config['crop_shape'],
)