def generate_data_additive_treatment_process(m, alpha=1, print_details=False):
normal_sigma_1 = dists.Normal(0, 1)
bern = dists.Bernoulli(0.5)
unif = dists.Uniform(0, 1)
z = normal_sigma_1.sample((m, ))
eps = normal_sigma_1.sample((m, ))
assert z.shape == eps.shape, (z.shape, eps.shape)
t = (z + eps) / 1.414
if print_details:
print("generated data with treated fraction {}".format(torch.mean(t)))
print("running with alpha = ", alpha)
y = 0.32 * normal_sigma_1.sample((m, )) + t + alpha * t * t * z
# replace 0*z with just z to use the confounder in the main code
return None, 0 * z, eps, t, y
def step(self, x, model):
sample = x.clone()
for i in range(self.dim):
lp_keep = model(sample).squeeze()
xi_keep = sample[:, i]
xi_change = 1. - xi_keep
sample_change = sample.clone()
sample_change[:, i] = xi_change
lp_change = model(sample_change).squeeze()
lp_update = lp_change - lp_keep
update_dist = dists.Bernoulli(logits=lp_update)
updates = update_dist.sample()
sample = sample_change * updates[:, None] + sample * (
1. - updates[:, None])
self.changes[i] = updates.mean()
return sample
def forward(self, x):
batch_sz = x.size()[0]
sz = self.q_pi_a.size()
p_pi = distributions.Beta(
torch.ones(sz) * self.p_pi_alpha,
torch.ones(sz) * self.p_pi_beta)
beta_a = F.softplus(self.q_pi_alpha) + 0.01
beta_b = F.softplus(self.q_pi_beta) + 0.01
q_pi = distributions.Beta(beta_a, beta_b)
# Differentiable Sample Knowles et al.
qpi_sample = q_pi.rsample()
q_z = shared.STRelaxedBernoulli(temperature=0.1, probs=qpi_sample)
z = q_z.rsample()
q_z = distributions.Bernoulli(probs=qpi_sample)
q_phi = distributions.Normal(loc=self.phi_mean,
scale=(self.phi_logvar / 2).exp())
q_w = distributions.Normal(loc=self.w_mean,
scale=(self.w_logvar / 2).exp())
# For now, just take the mean
phi = q_phi.mean
w = q_w.mean
# Alternatively, sample
# phi = q_phi.rsample()
# w = q_w.rsample()
# NLL
sinbasis = torch.ones(K, N_SAMPLES) * torch.arange(0, N_SAMPLES, 1)
for k in range(K):
sinbasis[k] = torch.sin(sinbasis[k] * phi[k])
x_mean = torch.mm(torch.mul(z, w),
sinbasis) # z and w multiplied elementwise
nll = -(distributions.Normal(loc=x_mean,
scale=self.sigma_n).log_prob(x))
return nll, p_pi, q_pi, q_z, q_phi, q_w, sinbasis
def __init__(self,
n_node,
avg_degree=2,
init_bias=0.,
learn_G=False,
learn_bias=False):
super().__init__()
g = ig.Graph.Erdos_Renyi(n_node, float(avg_degree) / float(n_node))
A = np.asarray(g.get_adjacency().data) # g.get_sparse_adjacency()
A = torch.tensor(A).float()
weights = torch.randn_like(A) * ((1. / avg_degree)**.5)
weights = weights * (1 - torch.tril(torch.ones_like(weights)))
weights = weights + weights.t()
self.G = nn.Parameter(A * weights, requires_grad=learn_G)
self.bias = nn.Parameter(torch.ones((n_node, )).float() * init_bias,
requires_grad=learn_bias)
self.init_dist = dists.Bernoulli(logits=2 * self.bias)
self.data_dim = n_node
def __init__(self,
in_channels=1,
out_dim=1,
probs_fn=torch.sigmoid,
sample_fn=lambda x: distributions.Bernoulli(probs=x).sample(),
n_gated=10,
gated_channels=128,
head_channels=32):
"""Initializes a new GatedPixelCNN instance.
Args:
in_channels: The number of channels in the input.
out_dim: The dimension of the output. Given input of the form NCHW, the
output from the GatedPixelCNN model will be N out_dim CHW.
probs_fn: See the base class.
sample_fn: See the base class.
n_gated: The number of gated layers (not including the input layers).
gated_channels: The number of channels to use in the gated layers.
head_channels: The number of channels to use in the 1x1 convolution blocks
in the head after all the gated channels.
"""
super().__init__(probs_fn, sample_fn)
self._out_dim = out_dim
self._input = GatedPixelCNNLayer(in_channels=in_channels,
out_channels=gated_channels,
kernel_size=7,
is_causal=True)
self._gated_layers = nn.ModuleList([
GatedPixelCNNLayer(in_channels=gated_channels,
out_channels=gated_channels,
kernel_size=3,
is_causal=False) for _ in range(n_gated)
])
self._head = nn.Sequential(
nn.ReLU(),
nn.Conv2d(in_channels=gated_channels,
out_channels=head_channels,
kernel_size=1), nn.ReLU(),
nn.Conv2d(in_channels=head_channels,
out_channels=self._out_dim * in_channels,
kernel_size=1))
def __init__(self,
in_channels=1,
out_dim=1,
probs_fn=torch.sigmoid,
sample_fn=lambda x: distributions.Bernoulli(probs=x).sample(),
n_residual=15,
residual_channels=128,
head_channels=32):
"""Initializes a new PixelCNN instance.
Args:
in_channels: The number of channels in the input image (typically either
1 or 3 for black and white or color images respectively).
out_dim: The dimension of the output. Given input of the form NCHW, the
output from the model will be N out_dim CHW.
probs_fn: See the base class.
sample_fn: See the base class.
n_residual: The number of residual blocks.
residual_channels: The number of channels to use in the residual layers.
head_channels: The number of channels to use in the two 1x1 convolutional
layers at the head of the network.
"""
super().__init__(probs_fn, sample_fn)
self._out_dim = out_dim
self._input = pg_nn.MaskedConv2d(is_causal=True,
in_channels=in_channels,
out_channels=2 * residual_channels,
kernel_size=7,
padding=3)
self._masked_layers = nn.ModuleList([
MaskedResidualBlock(n_channels=2 * residual_channels)
for _ in range(n_residual)
])
self._head = nn.Sequential(
nn.ReLU(),
nn.Conv2d(in_channels=2 * residual_channels,
out_channels=head_channels,
kernel_size=1), nn.ReLU(),
nn.Conv2d(in_channels=head_channels,
out_channels=self._out_dim * in_channels,
kernel_size=1))
def __init__(self,
dim,
init_sigma=.15,
init_bias=0.,
learn_G=False,
learn_sigma=False,
learn_bias=False,
lattice_dim=2):
super().__init__()
g = ig.Graph.Lattice(dim=[dim] * lattice_dim,
circular=True) # Boundary conditions
A = np.asarray(g.get_adjacency().data) # g.get_sparse_adjacency()
self.G = nn.Parameter(torch.tensor(A).float(), requires_grad=learn_G)
self.sigma = nn.Parameter(torch.tensor(init_sigma).float(),
requires_grad=learn_sigma)
self.bias = nn.Parameter(torch.ones(
(dim**lattice_dim, )).float() * init_bias,
requires_grad=learn_bias)
self.init_dist = dists.Bernoulli(logits=2 * self.bias)
self.data_dim = dim**lattice_dim
def generate(self, n_samples, args):
samples_time_seq = []
# initialize model
c = torch.zeros(n_samples, self.C * self.H * self.W).to(args.device)
h_dec = torch.zeros(n_samples, self.lstm_size).to(args.device)
c_dec = torch.zeros_like(h_dec).to(args.device)
# run for the number of time steps
for t in range(self.time_steps):
z_sample = D.Normal(0, 1).sample(
(n_samples, self.z_size)).to(args.device)
h_dec, c_dec = self.decoder(z_sample, (h_dec, c_dec))
c = c + self.write(h_dec)
x = D.Bernoulli(
logits=c.view(n_samples, self.C, self.H, self.W)).probs
samples_time_seq.append(x)
return samples_time_seq
def forward(self, p, inputs, sample=True, **kwargs):
''' Performs a forward pass through the network.
Args:
p (tensor): sampled points
inputs (tensor): conditioning input
sample (bool): whether to sample for z
'''
# print(p.shape)
batch_size = p.size(0)
c = self.encoder(inputs)
logits = self.decoder(p, c)
crf_out = None
if self.crf != None:
crf_out = self.crf(logits, p)
#crf_out = self.crf(logits.unsqueeze(1),p.permute(0,2,1)).squeeze()
p_r = dist.Bernoulli(logits=crf_out) #logits)
# return logits,p_r,crf_out
return crf_out, p_r
def get_losses(z_mean, z_std, x_pred, x_std, x):
kld = kld_loss(z_mean, z_std)
x_loss = 0
pred_i = 0 #x_pred is much longer than x if x has categorical variables with more categories than 2
for i, mode in enumerate(x_mode):
if mode == 0:
x_loss += -dist.Normal(loc=x_pred[:, pred_i],
scale=x_std[:, pred_i]).log_prob(
x[:, i]).sum()
pred_i += 1
elif mode == 2:
x_loss += -dist.Bernoulli(logits=x_pred[:, pred_i]).log_prob(
x[:, i]).sum()
pred_i += 1
else:
x_loss += -dist.Categorical(logits=x_pred[:, pred_i:pred_i +
mode]).log_prob(
x[:, i]).sum()
pred_i += mode
return kld, x_loss
def select_action(self, obs):
# [B, P], [B, P, A], [B, P]
(q, pi, beta) = self.net(obs, rnncs=self.rnncs)
self.rnncs_ = self.net.get_rnncs()
options_onehot = F.one_hot(self.options,
self.options_num).float() # [B, P]
options_onehot_expanded = options_onehot.unsqueeze(-1) # [B, P, 1]
pi = (pi * options_onehot_expanded).sum(-2) # [B, A]
if self.is_continuous:
mu = pi # [B, A]
log_std = self.log_std[self.options] # [B, A]
dist = td.Independent(td.Normal(mu, log_std.exp()), 1)
action = dist.sample().clamp(-1, 1) # [B, A]
log_prob = dist.log_prob(action).unsqueeze(-1) # [B, 1]
else:
logits = pi # [B, A]
norm_dist = td.Categorical(logits=logits)
action = norm_dist.sample() # [B,]
log_prob = norm_dist.log_prob(action).unsqueeze(-1) # [B, 1]
value = q_o = (q * options_onehot).sum(-1, keepdim=True) # [B, 1]
beta_adv = q_o - ((1 - self.eps) * q.max(-1, keepdim=True)[0] +
self.eps * q.mean(-1, keepdim=True)) # [B, 1]
max_options = q.argmax(-1) # [B, P] => [B, ]
beta_probs = (beta * options_onehot).sum(-1) # [B, P] => [B,]
beta_dist = td.Bernoulli(probs=beta_probs)
# <1 则不改变op， =1 则改变op
new_options = th.where(beta_dist.sample() < 1, self.options,
max_options)
self.new_options = th.where(self._done_mask, max_options, new_options)
self.oc_mask = (self.new_options == self.options).float()
acts_info = Data(
action=action,
value=value,
log_prob=log_prob + th.finfo().eps,
beta_advantage=beta_adv + self.dc,
last_options=self.options,
options=self.new_options,
reward_offset=-((1 - self.oc_mask) * self.dc).unsqueeze(-1))
if self.use_rnn:
acts_info.update(rnncs=self.rnncs)
return action, acts_info
def bernoulli_loss(reconstructions, images, activation, subtract_entropy=False):
"""Bernoulli loss"""
flattened_dim = images.size(1)*images.size(2)*images.size(3)
reconstructions = reconstructions.view(-1, flattened_dim)
images = images.view(-1, flattened_dim)
if subtract_entropy:
dist = distributions.Bernoulli(probs=torch.clamp(images, 1e-6, (1-1e-6)))
lower_bound = dist.entropy().sum(dim=1)
else:
lower_bound = 0
if activation == "logits":
loss = F.binary_cross_entropy_with_logits(reconstructions, images, reduction="none").sum(dim=1)
elif activation == "tanh":
reconstructions = torch.clamp(torch.tanh(reconstructions).div(2).add(0.5), 1e-6, (1-1e-6))
loss = -torch.sum(images.mul(torch.log(reconstructions)) + (1-images).mul(torch.log(1 - reconstructions)), dim=1)
else:
raise ValueError("unknown activation function")
return loss - lower_bound
def score_image(self, ctoken, image):
"""
Compute the log-likelihood of a binary image
Note: Should only be called from Model
Parameters
----------
image : (H,W) tensor
binary image to score
Returns
-------
ll : tensor
scalar; log-likelihood of the image
"""
pimg = self.get_pimg(ctoken)
bern = dist.Bernoulli(pimg)
ll = bern.log_prob(image)
ll = torch.sum(ll)
return ll
def __init__(
self,
in_channels=1,
out_dim=1,
probs_fn=torch.sigmoid,
sample_fn=lambda x: distributions.Bernoulli(probs=x).sample()):
"""Initializes a new TinyCNN instance.
Args:
in_channels: Number of input channels.
out_dim: Dimension of the output per channel.
probs_fn: See the base class.
sample_fn: See the base class.
"""
super().__init__(probs_fn, sample_fn)
self._out_dim = out_dim
self._conv = pg_nn.MaskedConv2d(is_causal=True,
in_channels=in_channels,
out_channels=out_dim * in_channels,
kernel_size=3,
padding=1)
def forward(self, x, mask):
"""
x.size() = ("seq", "batch", ...)
"""
# Compute the parameters for the Bernoulli
embeddings = self.compute_embedding(x)
dist_params = self.predict_distribution(embeddings)
dist_params = dist_params.squeeze()
# Sample
sampler = dist.Bernoulli(probs=dist_params)
actions = sampler.sample()
# Compute LogProba
log_probas = sampler.log_prob(actions)
log_probas = apply_mask(log_probas, mask)
# Compute Entropy
entropy = sampler.entropy()
entropy = apply_mask(log_probas, mask)
return actions, log_probas, entropy, dist_params
def _forward(self, x):
"""Computes the forward pass and samples a new output.
Returns:
(p_hat, x_hat) where p_hat is the probability distribution over dimensions
and x_hat is sampled from p_hat.
"""
# If the input is an image, flatten it during the forward pass.
original_shape = x.shape
if len(x.shape) > 2:
x = x.view(original_shape[0], -1)
in_W, in_b = self.params["in_W"], self.params["in_b"]
h_W, h_b = self.params["h_W"], self.params["h_b"]
batch_size = 1 if x is None else x.shape[0]
p_hat = []
x_hat = []
# Only the bias is used to compute the first hidden unit so we must
# replicate it to account for the batch size.
a = in_b.expand(batch_size, -1)
for i in range(self._input_dim):
h = torch.relu(a)
p_i = torch.sigmoid(h_b[i : i + 1] + h @ h_W[i : i + 1, :].t())
p_hat.append(p_i)
# Sample 'x' at dimension 'i' if it is not given.
x_i = x[:, i : i + 1]
x_i = torch.where(x_i < 0, distributions.Bernoulli(probs=p_i).sample(), x_i)
x_hat.append(x_i)
# We do not need to add in_b[i:i+1] when computing the other hidden units
# since it was already added when computing the first hidden unit.
a = a + x_i @ in_W[:, i : i + 1].t()
if x_hat:
return (
torch.cat(p_hat, dim=1).view(original_shape),
torch.cat(x_hat, dim=1).view(original_shape),
)
return []
def forward(self, encoder_output):
# encoder_output is a tensor of size [batch_size, max_length, input_embed]
W = self._W
logits = torch.einsum('ijk, kn, imn->ijm', encoder_output, W, encoder_output) # Readability
self.logit_bias = self._l
if self.use_bias: # Bias to control sparsity/density
logits += self.logit_bias
self.adj_prob = logits
self.samples = []
self.mask = 0
self.mask_scores = []
self.entropy = []
for i in range(self.max_length):
position = torch.ones([encoder_output.shape[0]],
device=self.device) * i
position = position.long()
# Update mask
self.mask = torch.zeros((encoder_output.shape[0], self.max_length),
device=self.device).scatter_(1, position.view(encoder_output.shape[0], 1), 1)
masked_score = self.adj_prob[:,i,:] - 100000000.*self.mask
prob = distr.Bernoulli(logits=masked_score) # probs input probability, logit input log_probability
sampled_arr = prob.sample() # Batch_size, seqlenght for just one node
sampled_arr.requires_grad=True
self.samples.append(sampled_arr)
self.mask_scores.append(masked_score)
self.entropy.append(prob.entropy())
return self.samples, self.mask_scores, self.entropy
def generate_image_data(num_samples, zdim, generator, t_a, t_b, y_a0, y_b0,
y_a1, y_b1, c_x, s_x):
z = torch.randn(num_samples, zdim)
with torch.no_grad():
image_expectations = (1 + generator(z[:, :, None, None])) / 2
images = dist.Bernoulli(
image_expectations).sample() #This way binary data
#image_means = generator(z[:,:,None,None])
#images = image_means# + torch.randn_like(image_means)*0.05 # <- this way continuous data
z_temp = z[:, 0][:, None].detach().numpy(
) #Use the first dimension for prediction of ordinary variables
x = torch.Tensor(
np.random.normal(
np.tile(c_x, (num_samples, 1)) * z_temp,
np.tile(s_x, (num_samples, 1)), (num_samples, 1)))
t = (np.random.random(
(num_samples, 1)) < sigmoid(t_a * z_temp + t_b)).astype(int)
y = torch.Tensor((np.random.random((num_samples, 1)) < sigmoid(y_a1*z_temp + y_b1)).astype(int)*t \
+ (np.random.random((num_samples, 1)) < sigmoid(y_a0*z_temp + y_b0)).astype(int)*(1-t))
t = torch.Tensor(t)
dataset = ImageDataset(images, x, t, y, zdim)
return z, images, x, t, y, dataset
def __init__(self,
in_channels,
in_size,
out_dim=1,
probs_fn=torch.sigmoid,
sample_fn=lambda x: distributions.Bernoulli(probs=x).sample(),
n_transformer_blocks=8,
n_attention_heads=4,
n_embedding_channels=16):
"""Initializes a new ImageGPT instance.
Args:
in_channels: The number of input channels.
in_size: Size of the input images. Used to create positional encodings.
out_dim: The dimension of the output. Given input of the form NCHW, the
output from the model will be N out_dim CHW.
probs_fn: See the base class.
sample_fn: See the base class.
n_transformer_blocks: Number of TransformerBlocks to use.
n_attention_heads: Number of attention heads to use.
n_embedding_channels: Number of attention embedding channels to use.
"""
super().__init__(probs_fn, sample_fn)
self._out_dim = out_dim
self._pos = nn.Parameter(torch.zeros(1, in_channels, in_size, in_size))
self._input = pg_nn.MaskedConv2d(
is_causal=True,
in_channels=in_channels,
out_channels=n_embedding_channels,
kernel_size=3,
padding=1)
self._transformer = nn.ModuleList(
TransformerBlock(n_channels=n_embedding_channels,
n_attention_heads=n_attention_heads)
for _ in range(n_transformer_blocks))
self._out = nn.Conv2d(in_channels=n_embedding_channels,
out_channels=self._out_dim * in_channels,
kernel_size=1)
def sample(self, n: int, device: str = 'cpu'):
"""
Sample n samples from the model.
"""
with torch.no_grad():
x = torch.zeros(n, self.input_size).to(device)
alphas = self.sample_alphas(n, device)
for idx in range(self.input_size):
alpha_row, alpha_column = alphas[:, idx], alphas[:, idx + 1]
theta = self.net(x).view(-1, self.input_size,
self.distribution_param_count,
*(self.alpha_dim, ) * 2)
batch_idx = torch.arange(theta.shape[0]).to(theta.device)
if self.distribution == 'binary':
beta_logits = theta[batch_idx, idx, 0, alpha_row,
alpha_column]
d = D.Bernoulli(logits=beta_logits)
x[:, idx] = d.sample()
elif self.distribution == 'categorical':
category_logits = theta[batch_idx, idx, :, alpha_row,
alpha_column]
d = D.Categorical(logits=category_logits)
x[:,
idx] = d.sample().float() / (self.kwargs['categories'] -
1)
elif self.distribution == 'gaussian':
mu = theta[batch_idx, idx, 0, alpha_row, alpha_column]
std = theta[batch_idx, idx, 1, alpha_row,
alpha_column].exp()
d = self._sigmoid_gaussian(mu, std)
x[:, idx] = d.sample()
assert (x >= 0).all()
assert (x <= 1).all()
return x, theta
def __init__(self,
*args,
word_dim=300,
hidden_dim=600,
word_dropout=0.0,
word_emb_dropout=0.0,
word_encoder=AbstractWordEncoder,
dst_encoder=AbstractDSTEncoder,
**kwargs):
super().__init__(*args, **kwargs)
self.word_dim = word_dim
self.hidden_dim = hidden_dim
self.word_dropout = word_dropout
self.word_emb_dropout = word_emb_dropout
self.word_encoder_cls = word_encoder
self.dst_encoder_cls = dst_encoder
self.word_encoder = self.word_encoder_cls(vocab=self.vocabs.word,
word_dim=self.word_dim)
self.sent_encoder = self.dst_encoder_cls(
ontology=self.vocabs.speaker_state["user"],
input_dim=self.word_dim,
hidden_dim=self.hidden_dim)
self.act_encoder = self.dst_encoder_cls(
ontology=self.vocabs.speaker_state["user"],
input_dim=self.word_dim,
hidden_dim=self.hidden_dim)
self.ont_encoder = self.dst_encoder_cls(
ontology=self.vocabs.speaker_state["user"],
input_dim=self.word_dim,
hidden_dim=self.hidden_dim)
self.sent_linear = nn.Linear(self.hidden_dim, 1)
self.act_weight = nn.Parameter(torch.tensor(0.5))
if self.word_dropout:
self.word_dropout_dist = dist.Bernoulli(self.word_dropout)
else:
self.word_dropout_dist = None
self.word_emb_dropout_layer = nn.Dropout(self.word_emb_dropout)
def decode(self, p, z, f3, f2, f1, **kwargs):
''' Returns occupancy probabilities for the sampled points.
Args:
p (tensor): points
z (tensor): latent code z
c (tensor): latent conditioned code c
'''
logits3 = self.decoder3(p, z, f3, **kwargs)
logits2 = self.decoder2(p, z, f2, **kwargs)
logits1 = self.decoder1(p, z, f1, **kwargs)
if self.local_feature_mask:
logits = logits3
w = (torch.abs(logits) < 1.5).detach().float()
logits = logits + self.logits2_ratio * logits2 * w
w = (torch.abs(logits) < 1).detach().float()
logits = logits + self.logits1_ratio * logits1 * w
else:
logits = logits3 + self.logits2_ratio * logits2 + self.logits1_ratio * logits1
p_r = dist.Bernoulli(logits=logits)
return p_r
def log_prob(self, x):
"""
Evaluate the log likelihood of a batch of samples.
"""
theta = self.net(x).view(-1, self.input_size,
self.distribution_param_count,
*(self.alpha_dim, ) * 2)
if self.distribution == 'binary':
d = D.Bernoulli(logits=theta[:, :, 0])
elif self.distribution == 'categorical':
d = D.Categorical(
logits=theta.permute(0, 1, 3, 4, 2)
) # We need to permute to put the categories on the last dimension for D.Categorical
elif self.distribution == 'gaussian':
mu, std = theta[:, :, 0], theta[:, :, 1].exp()
d = self._sigmoid_gaussian(mu, std)
# We need the indices for categorical distribution
if self.distribution == 'categorical':
x = x * (self.kwargs['categories'] - 1)
# Conditional probability matrices log p(x_i | alpha_{i-1}, alpha_i, x_{<i})
log_px_matrices = d.log_prob(x.unsqueeze(-1).unsqueeze(-1))
if self.alpha_dim > 1 or True:
# Multiply with the alpha marginals so that matrix multiplication corresponds to summing out alpha
joint_matrices = log_px_matrices[:, :-1] + self._log_p_alpha()
# Reduce all of the "inner" matrices (all except x_1 and x_n)
left_reduce_product = reduce(fast_logexpmv,
joint_matrices.transpose(0, 1))
p_xn_g_alpha_n1 = log_px_matrices[:, -1, :, 0:1]
log_px = fast_logexpmv(left_reduce_product, p_xn_g_alpha_n1)
else:
log_px = log_px_matrices.sum(1).squeeze(-1).squeeze(-1)
return log_px, theta
def mixture(self, output):
e_est, mu_x, mu_y, std_x_est, std_y_est, rho_est, pi_est = output.squeeze(
).requires_grad_(True)
# From estimated to real value
e = torch.sigmoid(e_est).requires_grad_(True)
rho = torch.tanh(rho_est).requires_grad_(True)
std_x = torch.exp(std_x_est).requires_grad_(True)
std_y = torch.exp(std_y_est).requires_grad_(True)
cov = (rho * std_x * std_y).requires_grad_(True)
# Parameters of the bivariate Gaussian
cov_matrix = torch.tensor([[std_x**2, cov],
[cov, std_y**2]]).requires_grad_(True)
means = torch.tensor([mu_x, mu_y]).requires_grad_(True)
# Instantiate distributions
bivariate_distrib = distrib.MultivariateNormal(means, cov_matrix)
bernoulli_distrib = distrib.Bernoulli(torch.tensor([e]))
# Sample from distributions
x, y = bivariate_distrib.sample().requires_grad_(True)
end_of_stroke = bernoulli_distrib.sample().requires_grad_(True)
return torch.tensor([end_of_stroke, x, y]).requires_grad_(True)
def extra_pass(self, name, passes):
# Render image with no NEE
if name[:5] == "nonee":
samples = int(name[5:])
beauty = self.renderer.get_image(samples, nee=False)["beauty"]
beauty = torch.from_numpy(beauty).float().to(self.device).permute(
2, 0, 1)
beauty = beauty.reshape((1, 1, *beauty.shape))
return beauty
# Create a noisy image
if name == "noise":
strength = 0.25
tmp = passes["beauty"].clone()
noise_level = torch.ones_like(tmp)[:, :, 0]
noise_level -= strength
snp_noise = D.Bernoulli(probs=noise_level).sample()
tmp = D.Normal(tmp, 2 * strength).sample()
tmp = snp_noise * tmp
return tmp
# Create a gray scale image
if name == "gray":
tmp = passes["beauty"].clone()
tmp = torch.mean(tmp, dim=2, keepdim=True)
return tmp
def crossover(parents, fitness, population):
"""
Crossover parents to get next generation. Parents are coupled and the features are sampled
according to parents' fitness scores.
Args:
parents (1D array): indexes of sampled parents, [2 * (population_size - 1)]
fitness (1D array): population fitness scores, [population_size]
population (4D array): current generation features, [population_size x n_channels x h x w]
Returns:
children (4D array): next generation features, [population_size - 1 x n_features]
"""
_, nchannels, h, w = population.shape
fitness_pairs = fitness[parents.long()].view(-1, 2)
prob = fitness_pairs[:, 0] / fitness_pairs.sum(1)
parental_bernoulli = td.Bernoulli(prob)
inherit_mask = parental_bernoulli.sample_n(nchannels * h *
w) # [N-1, nchannels * h * w]
inherit_mask = inherit_mask.view(-1, nchannels, h, w)
parent_features = population[parents.long()]
children = torch.cuda.FloatTensor(inherit_mask.shape)
children = where(inherit_mask, parent_features[::2], parent_features[1::2])
return children
def compute_loss(p, occ, images):
''' Computes the loss.
'''
inputs = images.to(device)
kwargs = {}
p0_z = dist.Normal(torch.tensor([]).to(device), torch.tensor([]).to(device))
c = encode_inputs(inputs)
q_z = infer_z(p, occ, c, **kwargs)
z = q_z.rsample()
# KL-divergence
kl = dist.kl_divergence(q_z, p0_z).sum(dim=-1)
loss = kl.mean()
# General points
logits = dist.Bernoulli(logits=occ_net.loss_addon(p, z, c).squeeze()).logits
# logits = logits.usqueeze(-11)
loss_i = F.binary_cross_entropy_with_logits(
logits, occ, reduction='none')
loss = loss + loss_i.sum(-1).mean()
return loss
def _sample(self):
sample_shape = self.local_rate.shape
locations = D.Bernoulli(self.local_rate).sample()
zero_point_five = torch.tensor(0.5, device=self.device)
x_offset = D.Uniform(
low=0 - zero_point_five,
high=0 + zero_point_five,
).sample(sample_shape=sample_shape)
y_offset = D.Uniform(
low=0 - zero_point_five,
high=0 + zero_point_five,
).sample(sample_shape=sample_shape)
z_offset = D.Uniform(
low=0 - zero_point_five,
high=0 + zero_point_five,
).sample(sample_shape=sample_shape)
intensities = D.Uniform(low=self.min_int,
high=1.0).sample(sample_shape=sample_shape)
x_offset *= locations
y_offset *= locations
z_offset *= locations
intensities *= locations
return locations, x_offset, y_offset, z_offset, intensities
def __init__(self, input_size, hidden_size, use_bias=True,noisin_noise_type=None,noisin_noise_parama=0.0,noisin_noise_paramb=0.0):
"""
Most parts are copied from torch.nn.LSTMCell.
"""
super(NoisinCell, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.use_bias = use_bias
self.weight_ih = nn.Parameter(
torch.Tensor(4 * hidden_size,input_size))
self.weight_hh = nn.Parameter(
torch.Tensor(4 * hidden_size,hidden_size))
self.weight_hh_wdrop = None
if use_bias:
self.bias = nn.Parameter(torch.Tensor(4 * hidden_size))
else:
self.register_parameter('bias', None)
self.reset_parameters()
self.noisin_noise_type = noisin_noise_type
self.noisin_noise_parama = noisin_noise_parama
self.noisin_noise_paramb = noisin_noise_paramb
self.bernoulli_dist = tdist.Bernoulli(noisin_noise_parama)
self.gamma_dist = tdist.Gamma(noisin_noise_parama,noisin_noise_paramb) # alpha, gammma
def get_mnist_loaders(batch_size,
dynamically_binarize=False,
resize_to_32=False):
"""Create train and test loaders for the MNIST dataset.
Args:
batch_size: The batch size to use.
dynamically_binarize: Whether to dynamically binarize images values to {0, 1}.
resize_to_32: Whether to resize the images to 32x32.
Returns:
Tuple of (train_loader, test_loader).
"""
transform = [transforms.ToTensor()]
if dynamically_binarize:
transform.append(lambda x: distributions.Bernoulli(probs=x).sample())
if resize_to_32:
transform.append(lambda x: F.pad(x, (2, 2, 2, 2)))
transform = transforms.Compose(transform)
train_loader = data.DataLoader(
datasets.MNIST("/tmp/data",
train=True,
download=True,
transform=transform),
batch_size=batch_size,
shuffle=True,
num_workers=8,
)
test_loader = data.DataLoader(
datasets.MNIST("/tmp/data",
train=False,
download=True,
transform=transform),
batch_size=batch_size,
num_workers=8,
)
return train_loader, test_loader