def __init__(self, in_channels, feat_dim, nf_particle, nf_effect, g_dim, I_factor=10, psteps=1, n_timesteps=1, ngf=8, image_size=[64, 64]): super().__init__() out_channels = 1 n_layers = int(np.log2(image_size[0])) - 1 # Positional encoding buffers x = torch.linspace(-1, 1, image_size[0]) y = torch.linspace(-1, 1, image_size[1]) x_grid, y_grid = torch.meshgrid(x, y) # Add as constant, with extra dims for N and C self.register_buffer('x_grid_enc', x_grid.view((1, 1, 1) + x_grid.shape)) self.register_buffer('y_grid_enc', y_grid.view((1, 1, 1) + y_grid.shape)) # Temporal encoding buffers if n_timesteps > 1: t = torch.linspace(-1, 1, n_timesteps) # Add as constant, with extra dims for N and C self.register_buffer('t_grid', t) # Set state dim with config, depending on how many time-steps we want to take into account self.image_size = image_size self.n_timesteps = n_timesteps self.state_dim = feat_dim self.I_factor = I_factor self.psteps = psteps self.g_dim = g_dim # self.att_image_encoder = AttImageEncoder(in_channels, feat_dim, ngf, n_layers) self.image_encoder = ImageEncoder(in_channels, feat_dim * 4, ngf, n_layers) # feat_dim * 2 if sample here self.image_decoder = ImageDecoder(g_dim, out_channels, ngf, n_layers) # self.image_decoder = SimpleSBImageDecoder(feat_dim, out_channels, ngf, n_layers, image_size) self.koopman = KoopmanOperators(self.state_dim * 4, nf_particle, nf_effect, g_dim, n_timesteps)
def __init__(self, in_channels, feat_dim, nf_particle, nf_effect, g_dim, u_dim, n_objects, I_factor=10, n_blocks=1, psteps=1, n_timesteps=1, ngf=8, image_size=[64, 64]): super().__init__() out_channels = 1 n_layers = int(np.log2(image_size[0])) - 1 self.u_dim = u_dim # Temporal encoding buffers if n_timesteps > 1: t = torch.linspace(-1, 1, n_timesteps) # Add as constant, with extra dims for N and C self.register_buffer('t_grid', t) # Set state dim with config, depending on how many time-steps we want to take into account self.image_size = image_size self.n_timesteps = n_timesteps self.state_dim = feat_dim self.I_factor = I_factor self.psteps = psteps self.g_dim = g_dim self.n_objects = n_objects self.softmax = nn.Softmax(dim=-1) self.sigmoid = nn.Sigmoid() self.linear_g = nn.Linear(g_dim, g_dim) self.linear_u = nn.Linear(g_dim, u_dim) self.initial_conditions = nn.Sequential(nn.Linear(feat_dim * n_timesteps * 2, feat_dim * n_timesteps), nn.ReLU(), nn.Linear(feat_dim * n_timesteps, g_dim * 2)) self.image_encoder = ImageEncoder(in_channels, feat_dim * 2, n_objects, ngf, n_layers) # feat_dim * 2 if sample here self.image_decoder = ImageDecoder(g_dim, out_channels, ngf, n_layers) self.koopman = KoopmanOperators(feat_dim * 2, nf_particle * 2, nf_effect * 2, g_dim, u_dim, n_timesteps, n_blocks)
def __init__(self, in_channels, feat_dim, nf_particle, nf_effect, g_dim, u_dim, n_objects, I_factor=10, n_blocks=1, psteps=1, n_timesteps=1, ngf=8, image_size=[64, 64]): super().__init__() out_channels = 1 n_layers = int(np.log2(image_size[0])) - 1 self.u_dim = u_dim # Temporal encoding buffers if n_timesteps > 1: t = torch.linspace(-1, 1, n_timesteps) # Add as constant, with extra dims for N and C self.register_buffer('t_grid', t) # Set state dim with config, depending on how many time-steps we want to take into account self.image_size = image_size self.n_timesteps = n_timesteps self.state_dim = feat_dim self.I_factor = I_factor self.psteps = psteps self.g_dim = g_dim self.n_objects = n_objects self.softmax = nn.Softmax(dim=-1) self.sigmoid = nn.Sigmoid() # self.linear_g = nn.Linear(g_dim, g_dim) # self.initial_conditions = nn.Sequential(nn.Linear(feat_dim * n_timesteps * 2, feat_dim * n_timesteps), # nn.ReLU(), # nn.Linear(feat_dim * n_timesteps, g_dim * 2)) feat_dyn_dim = feat_dim // 6 self.feat_dyn_dim = feat_dyn_dim self.content = None self.reverse = False self.hankel = True self.ini_alpha = 1 self.incr_alpha = 0.5 # self.linear_u = nn.Linear(g_dim + u_dim, u_dim * 2) # self.linear_u_2_f = nn.Linear(feat_dyn_dim * n_timesteps, u_dim * 2) # self.linear_u_T_f = nn.Linear(feat_dyn_dim * n_timesteps, u_dim) if self.hankel: # self.linear_u_2_g = nn.Linear(g_dim * self.n_timesteps, u_dim * 2) # self.linear_u_2_g = nn.Sequential(nn.Linear(g_dim * self.n_timesteps, g_dim), # nn.ReLU(), # nn.Linear(g_dim, u_dim * 2)) # self.linear_u_all_g = nn.Sequential(nn.Linear(g_dim * self.n_timesteps, g_dim), # nn.ReLU(), # nn.Linear(g_dim, u_dim + 1)) # self.gru_u_all_g = nn.GRU(g_dim, u_dim + 1, num_layers = 2, batch_first=True) self.linear_u_1_g = nn.Sequential( nn.Linear(g_dim * self.n_timesteps, g_dim), nn.ReLU(), nn.Linear(g_dim, u_dim)) else: self.linear_u_2_g = nn.Linear(g_dim, u_dim * 2) self.linear_f_mu = nn.Linear(feat_dim * 2, feat_dim) self.linear_f_logvar = nn.Linear(feat_dim * 2, feat_dim) self.image_encoder = ImageEncoder( in_channels, feat_dim * 2, n_objects, ngf, n_layers) # feat_dim * 2 if sample here self.image_decoder = ImageDecoder(feat_dim, out_channels, ngf, n_layers) self.koopman = KoopmanOperators(feat_dyn_dim, nf_particle, nf_effect, g_dim, u_dim, n_timesteps, n_blocks)