def __init__(self,
num_steps,
x_size,
window_size,
z_what_size,
rnn_hidden_size,
encoder_net=[],
decoder_net=[],
predict_net=[],
embed_net=None,
bl_predict_net=[],
non_linearity='ReLU',
decoder_output_bias=None,
decoder_output_use_sigmoid=False,
use_masking=True,
use_baselines=True,
baseline_scalar=None,
scale_prior_mean=3.0,
scale_prior_sd=0.1,
pos_prior_mean=0.0,
pos_prior_sd=1.0,
likelihood_sd=0.3,
use_cuda=False):
super(AIR, self).__init__()
self.num_steps = num_steps
self.x_size = x_size
self.window_size = window_size
self.z_what_size = z_what_size
self.rnn_hidden_size = rnn_hidden_size
self.use_masking = use_masking
self.use_baselines = use_baselines
self.baseline_scalar = baseline_scalar
self.likelihood_sd = likelihood_sd
self.use_cuda = use_cuda
prototype = torch.tensor(0.).cuda() if use_cuda else torch.tensor(0.)
self.options = dict(dtype=prototype.dtype, device=prototype.device)
self.z_pres_size = 1
self.z_where_size = 3
# By making these parameters they will be moved to the gpu
# when necessary. (They are not registered with pyro for
# optimization.)
self.z_where_loc_prior = nn.Parameter(torch.FloatTensor(
[scale_prior_mean, pos_prior_mean, pos_prior_mean]),
requires_grad=False)
self.z_where_scale_prior = nn.Parameter(torch.FloatTensor(
[scale_prior_sd, pos_prior_sd, pos_prior_sd]),
requires_grad=False)
# Create nn modules.
rnn_input_size = x_size**2 if embed_net is None else embed_net[-1]
rnn_input_size += self.z_where_size + z_what_size + self.z_pres_size
nl = getattr(nn, non_linearity)
self.rnn = nn.LSTMCell(rnn_input_size, rnn_hidden_size)
self.encode = Encoder(window_size**2, encoder_net, z_what_size, nl)
self.decode = Decoder(window_size**2, decoder_net, z_what_size,
decoder_output_bias, decoder_output_use_sigmoid,
nl)
self.predict = Predict(rnn_hidden_size, predict_net, self.z_pres_size,
self.z_where_size, nl)
self.embed = Identity() if embed_net is None else MLP(
x_size**2, embed_net, nl, True)
self.bl_rnn = nn.LSTMCell(rnn_input_size, rnn_hidden_size)
self.bl_predict = MLP(rnn_hidden_size, bl_predict_net + [1], nl)
self.bl_embed = Identity() if embed_net is None else MLP(
x_size**2, embed_net, nl, True)
# Create parameters.
self.h_init = nn.Parameter(torch.zeros(1, rnn_hidden_size))
self.c_init = nn.Parameter(torch.zeros(1, rnn_hidden_size))
self.bl_h_init = nn.Parameter(torch.zeros(1, rnn_hidden_size))
self.bl_c_init = nn.Parameter(torch.zeros(1, rnn_hidden_size))
self.z_where_init = nn.Parameter(torch.zeros(1, self.z_where_size))
self.z_what_init = nn.Parameter(torch.zeros(1, self.z_what_size))
if use_cuda:
self.cuda()
def __init__(self,
num_steps,
x_size,
window_size,
z_what_size,
rnn_hidden_size,
encoder_net=[],
decoder_net=[],
predict_net=[],
embed_net=None,
bl_predict_net=[],
non_linearity='ReLU',
decoder_output_bias=None,
decoder_output_use_sigmoid=False,
use_masking=True,
use_baselines=True,
baseline_scalar=None,
fudge_z_pres=False,
use_cuda=False):
super(AIR, self).__init__()
self.num_steps = num_steps
self.x_size = x_size
self.window_size = window_size
self.z_what_size = z_what_size
self.rnn_hidden_size = rnn_hidden_size
self.use_masking = use_masking and not fudge_z_pres
self.use_baselines = use_baselines and not fudge_z_pres
self.baseline_scalar = baseline_scalar
self.fudge_z_pres = fudge_z_pres
self.use_cuda = use_cuda
self.z_pres_size = 1
self.z_where_size = 3
# By making these parameters they will be moved to the gpu
# when necessary. (They are not registered with pyro for
# optimization.)
self.z_where_mu_prior = nn.Parameter(torch.FloatTensor([3.0, 0, 0]),
requires_grad=False)
self.z_where_sigma_prior = nn.Parameter(torch.FloatTensor([0.1, 1, 1]),
requires_grad=False)
# Create nn modules.
rnn_input_size = x_size**2 if embed_net is None else embed_net[-1]
rnn_input_size += self.z_where_size + z_what_size + self.z_pres_size
nl = getattr(nn, non_linearity)
self.rnn = nn.LSTMCell(rnn_input_size, rnn_hidden_size)
self.encode = Encoder(window_size**2, encoder_net, z_what_size, nl)
self.decode = Decoder(window_size**2, decoder_net, z_what_size,
decoder_output_bias, decoder_output_use_sigmoid,
nl)
self.predict = Predict(rnn_hidden_size, predict_net, self.z_pres_size,
self.z_where_size, nl)
self.embed = Identity() if embed_net is None else MLP(
x_size**2, embed_net, nl, True)
self.bl_rnn = nn.LSTMCell(rnn_input_size, rnn_hidden_size)
self.bl_predict = MLP(rnn_hidden_size, bl_predict_net + [1], nl)
self.bl_embed = Identity() if embed_net is None else MLP(
x_size**2, embed_net, nl, True)
# Create parameters.
self.h_init = zeros(1, rnn_hidden_size)
self.c_init = zeros(1, rnn_hidden_size)
self.bl_h_init = zeros(1, rnn_hidden_size)
self.bl_c_init = zeros(1, rnn_hidden_size)
self.z_where_init = zeros(1, self.z_where_size)
self.z_what_init = zeros(1, self.z_what_size)
if use_cuda:
self.cuda()