-
Notifications
You must be signed in to change notification settings - Fork 1
/
utils.py
188 lines (159 loc) · 7.15 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
import torch
import torch.nn.functional as F
from torch import nn
import numpy as np
EPS = 1e-20
def sample_gumbel(shape, device):
U = torch.rand(shape)
U = U.to(device)
return -torch.log(-torch.log(U + EPS) + EPS)
def gumbel_softmax_sample(logits, temperature):
y = logits + sample_gumbel(logits.size(), logits.device)
return F.softmax(y / temperature, dim=-1)
def kl_categorical(q, k):
"""
Calculating KL between categorical q() and an uniform categorical p with k classes
"""
log_ratio = torch.log(q * k + EPS) # q / (1/k) = q*k
kl = torch.sum(q * log_ratio, dim=-1)
return kl
def sample_from_discretized_mix_logistic(l):
"""
Code taken from pytorch adaptation of original PixelCNN++ tf implementation
https://github.com/pclucas14/pixel-cnn-pp
"""
def to_one_hot(tensor, n):
one_hot = torch.zeros(tensor.size() + (n,))
one_hot = one_hot.to(tensor.device)
one_hot.scatter_(len(tensor.size()), tensor.unsqueeze(-1), 1.)
return one_hot
# Pytorch ordering
l = l.permute(0, 2, 3, 1)
ls = [int(y) for y in l.size()]
xs = ls[:-1] + [3]
# here and below: unpacking the params of the mixture of logistics
nr_mix = int(ls[-1] / 10)
# unpack parameters
logit_probs = l[:, :, :, :nr_mix]
l = l[:, :, :, nr_mix:].contiguous().view(xs + [nr_mix * 3])
# sample mixture indicator from softmax
temp = torch.FloatTensor(logit_probs.size())
if l.is_cuda:
temp = temp.cuda()
temp.uniform_(1e-5, 1. - 1e-5)
temp = logit_probs.data - torch.log(-torch.log(temp))
_, argmax = temp.max(dim=3)
one_hot = to_one_hot(argmax, nr_mix)
sel = one_hot.view(xs[:-1] + [1, nr_mix])
# select logistic parameters
means = torch.sum(l[:, :, :, :, :nr_mix] * sel, dim=4)
log_scales = torch.clamp(torch.sum(l[:, :, :, :, nr_mix:2 * nr_mix] * sel,
dim=4),
min=-7.)
coeffs = torch.sum(torch.tanh(l[:, :, :, :, 2 * nr_mix:3 * nr_mix]) * sel,
dim=4)
# sample from logistic & clip to interval
# we don't actually round to the nearest 8bit value when sampling
u = torch.FloatTensor(means.size())
if l.is_cuda:
u = u.cuda()
u.uniform_(1e-5, 1. - 1e-5)
u = nn.Parameter(u)
x = means + torch.exp(log_scales) * (torch.log(u) - torch.log(1. - u))
x0 = torch.clamp(torch.clamp(x[:, :, :, 0], min=-1.), max=1.)
x1 = torch.clamp(torch.clamp(x[:, :, :, 1] + coeffs[:, :, :, 0] * x0,
min=-1.),
max=1.)
x2 = torch.clamp(torch.clamp(x[:, :, :, 2] + coeffs[:, :, :, 1] * x0 +
coeffs[:, :, :, 2] * x1,
min=-1.),
max=1.)
out = torch.cat([
x0.view(xs[:-1] + [1]),
x1.view(xs[:-1] + [1]),
x2.view(xs[:-1] + [1])
],
dim=3)
# put back in Pytorch ordering
out = out.permute(0, 3, 1, 2)
return out
def discretized_mix_logistic_loss(x, l):
"""
log-likelihood for mixture of discretized logistics, assumes the data
has been rescaled to [-1,1] interval
Code taken from pytorch adaptation of original PixelCNN++ tf implementation
https://github.com/pclucas14/pixel-cnn-pp
"""
# channels last
x = x.permute(0, 2, 3, 1)
l = l.permute(0, 2, 3, 1)
# true image (i.e. labels) to regress to, e.g. (B,32,32,3)
xs = [int(y) for y in x.size()]
# predicted distribution, e.g. (B,32,32,100)
ls = [int(y) for y in l.size()]
# here and below: unpacking the params of the mixture of logistics
nr_mix = int(ls[-1] / 10)
logit_probs = l[:, :, :, :nr_mix]
l = l[:, :, :, nr_mix:].contiguous().view(
xs + [nr_mix * 3]) # 3 for mean, scale, coef
means = l[:, :, :, :, :nr_mix]
# log_scales = torch.max(l[:, :, :, :, nr_mix:2 * nr_mix], -7.)
log_scales = torch.clamp(l[:, :, :, :, nr_mix:2 * nr_mix], min=-7.)
coeffs = torch.tanh(l[:, :, :, :, 2 * nr_mix:3 * nr_mix])
# here and below: getting the means and adjusting them based on preceding
# sub-pixels
x = x.contiguous()
x = x.unsqueeze(-1) + nn.Parameter(torch.zeros(xs + [nr_mix]).to(x.device),
requires_grad=False)
m2 = (means[:, :, :, 1, :] + coeffs[:, :, :, 0, :] * x[:, :, :, 0, :]).view(
xs[0], xs[1], xs[2], 1, nr_mix)
m3 = (means[:, :, :, 2, :] + coeffs[:, :, :, 1, :] * x[:, :, :, 0, :] +
coeffs[:, :, :, 2, :] * x[:, :, :, 1, :]).view(
xs[0], xs[1], xs[2], 1, nr_mix)
means = torch.cat((means[:, :, :, 0, :].unsqueeze(3), m2, m3), dim=3)
centered_x = x - means
inv_stdv = torch.exp(-log_scales)
plus_in = inv_stdv * (centered_x + 1. / 255.)
cdf_plus = torch.sigmoid(plus_in)
min_in = inv_stdv * (centered_x - 1. / 255.)
cdf_min = torch.sigmoid(min_in)
# log probability for edge case of 0 (before scaling)
log_cdf_plus = plus_in - F.softplus(plus_in)
# log probability for edge case of 255 (before scaling)
log_one_minus_cdf_min = -F.softplus(min_in)
cdf_delta = cdf_plus - cdf_min # probability for all other cases
mid_in = inv_stdv * centered_x
# log probability in the center of the bin, to be used in extreme cases
# (not actually used in our code)
log_pdf_mid = mid_in - log_scales - 2. * F.softplus(mid_in)
# now select the right output: left edge case, right edge case, normal
# case, extremely low prob case (doesn't actually happen for us)
# this is what we are really doing, but using the robust version below
# for extreme cases in other applications and to avoid NaN issue with tf.select()
# log_probs = tf.select(x < -0.999, log_cdf_plus, tf.select(x > 0.999,
# log_one_minus_cdf_min, tf.log(cdf_delta)))
# robust version, that still works if probabilities are below 1e-5 (which
# never happens in our code)
# tensorflow backpropagates through tf.select() by multiplying with zero
# instead of selecting: this requires use to use some ugly tricks to avoid
# potential NaNs
# the 1e-12 in tf.maximum(cdf_delta, 1e-12) is never actually used as
# output, it's purely there to get around the tf.select() gradient issue
# if the probability on a sub-pixel is below 1e-5, we use an approximation
# based on the assumption that the log-density is constant in the bin of
# the observed sub-pixel value
inner_inner_cond = (cdf_delta > 1e-5).float()
inner_inner_out = inner_inner_cond * torch.log(
torch.clamp(cdf_delta, min=1e-12)) + (1. - inner_inner_cond) * (
log_pdf_mid - np.log(127.5))
inner_cond = (x > 0.999).float()
inner_out = inner_cond * log_one_minus_cdf_min + (
1. - inner_cond) * inner_inner_out
cond = (x < -0.999).float()
log_probs = cond * log_cdf_plus + (1. - cond) * inner_out
log_probs = torch.sum(log_probs, dim=3) + torch.log_softmax(logit_probs,
dim=-1)
log_probs = torch.logsumexp(log_probs, dim=-1)
# return -torch.sum(log_probs)
loss_sep = -log_probs.sum((1, 2)) # keep batch dimension
return loss_sep