def _gru(x, h, w, b, with_bias):
"""GRU cell.
Args:
x (:obj:`~nnabla.Variable`): Input data.
h (:obj:`~nnabla.Variable`): Hidden state.
w (:obj:`~nnabla.Variable`): Weight.
b (:obj:`~nnabla.Variable`): Bias.
with_bias (bool): Include the bias or not.
"""
hidden_size = h.shape[1]
xh = F.concatenate(*(x, h), axis=1)
w0, w1, w2 = F.split(w, axis=0)
b0 = b1 = b2 = b3 = None
if with_bias:
b0, b1, b2, b3 = F.split(b, axis=0)
r_t = F.sigmoid(F.affine(xh, F.transpose(w0, (1, 0)), b0))
z_t = F.sigmoid(F.affine(xh, F.transpose(w1, (1, 0)), b1))
w2_0 = w2[:, :w2.shape[1] - hidden_size]
w2_1 = w2[:, w2.shape[1] - hidden_size:]
n_t = F.tanh(
F.affine(x, F.transpose(w2_0, (1, 0)), b2) +
r_t * F.affine(h, F.transpose(w2_1, (1, 0)), b3))
h_t = (1 - z_t) * n_t + z_t * h
return h_t
def _lstm_cell(self, name, n_hidden, x_in, h=None, c=None):
if h is None:
h = nn.Variable.from_numpy_array(
np.zeros((self._batch_size, self._cols_size)))
if c is None:
c = nn.Variable.from_numpy_array(
np.zeros((self._batch_size, n_hidden)))
h = F.concatenate(h, x_in, axis=1) # LSTM_Concatenate -> cols_size * 2
with nn.parameter_scope(name + '_Affine'): # LSTM_Affine -> n_hidden
h1 = PF.affine(h, (n_hidden, ), base_axis=1)
with nn.parameter_scope(name + '_IGate'): # LSTM_IGate -> n_hidden
h2 = PF.affine(h, (n_hidden, ), base_axis=1)
with nn.parameter_scope(name + '_FGate'): # LSTM_FGate -> n_hidden
h3 = PF.affine(h, (n_hidden, ), base_axis=1)
with nn.parameter_scope(name + '_OGate'): # LSTM_OGate -> n_hidden
h4 = PF.affine(h, (n_hidden, ), base_axis=1)
h1 = F.tanh(h1) # LSTM_Tanh
h2 = F.sigmoid(h2) # LSTM_Sigmoid
h3 = F.sigmoid(h3) # LSTM_Sigmoid_2
h4 = F.sigmoid(h4) # LSTM_Sigmoid_3
h5 = F.mul2(h2, h1) # LSTM_Mul2 -> n_hidden
h6 = F.mul2(h3, c) # LSTM_Mul2_2 -> n_hidden
h7 = F.add2(h5, h6, inplace=True) # LSTM_Add2 -> n_hidden
h8 = F.tanh(h7) # LSTM_Tanh_2 -> n_hidden
h9 = F.mul2(h4, h8) # LSTM_Mul2_3 -> n_hidden
c = h7 # LSTM_C
h = h9 # LSTM_H
return (h, c)
def LSTMCell(x, h2, h1):
units = h1.shape[1]
#first stack h2=hidden, h1= cell
h2 = F.concatenate(h2, x, axis=1)
h3 = PF.affine(h2, (units), name='Affine')
h4 = PF.affine(h2, (units), name='InputGate')
h5 = PF.affine(h2, (units), name='ForgetGate')
h6 = PF.affine(h2, (units), name='OutputGate')
h3 = F.tanh(h3)
h4 = F.sigmoid(h4)
h5 = F.sigmoid(h5)
h6 = F.sigmoid(h6)
h4 = F.mul2(h4, h3)
h5 = F.mul2(h5, h1)
h4 = F.add2(h4, h5, True)
h7 = F.tanh(h4)
h6 = F.mul2(h6, h7)
return h6, h4 # hidden, cell
def conv_lstm_cell(input_tensor, cur_state, n_filt, kernel_size):
"""
conv lstm cell definition
"""
def split(inp):
_, channels, _, _ = inp.shape
channels = channels / 4
return inp[:, :channels, :, :], inp[:, channels:2 * channels, :, :], \
inp[:, 2 * channels:3 * channels, :, :], \
inp[:, 3 * channels:4 * channels, :, :]
h_cur, c_cur = cur_state
# concatenate along channel axis
combined = F.concatenate(*[input_tensor, h_cur], axis=1)
combined_conv = conv2d(combined, 4 * n_filt, kernel_size, 1, kernel_size // 2,
name='conv_lstm_cell')
cc_i, cc_f, cc_o, cc_g = split(combined_conv)
act_i = F.sigmoid(cc_i)
act_f = F.sigmoid(cc_f)
act_o = F.sigmoid(cc_o)
act_g = F.tanh(cc_g)
c_next = F.add2(act_f * c_cur, act_i * act_g)
h_next = act_o * F.tanh(c_next)
return h_next, c_next
def yolov2_activate(x, anchors, biases):
shape = x.shape
y = F.reshape(x, (
shape[0],
anchors,
-1,
) + shape[2:])
stop = list(y.shape)
stop[2] = 2
t_xy = F.slice(y, (0, 0, 0, 0, 0), stop)
stop[2] = 4
t_wh = F.slice(y, (0, 0, 2, 0, 0), stop)
stop[2] = 5
t_o = F.slice(y, (0, 0, 4, 0, 0), stop)
stop[2] = y.shape[2]
t_p = F.slice(y, (0, 0, 5, 0, 0), stop)
t_xy = F.sigmoid(t_xy)
t_wh = F.exp(t_wh)
t_o = F.sigmoid(t_o)
t_p = F.softmax(t_p, axis=2)
t_x, t_y, t_wh = yolov2_image_coordinate(t_xy, t_wh, biases)
y = F.concatenate(t_x, t_y, t_wh, t_o, t_p, axis=2)
y = F.transpose(y, (0, 1, 3, 4, 2)).reshape(
(shape[0], -1, shape[1] / anchors))
return y
def gated_conv(self,
x,
kernel_shape,
h=None,
mask_type='',
gated=True,
payload=None,
return_payload=False,
scope_name='gated_conv'):
pad_dim_0 = (kernel_shape[0] - 1) / 2
pad_dim_1 = (kernel_shape[1] - 1) / 2
if mask_type == '':
mask_type = self.mask_type_B
with nn.parameter_scope(scope_name):
if gated:
out_f = PF.convolution(x,
self.num_features,
kernel_shape,
pad=(pad_dim_0, pad_dim_1),
apply_w=mask_type,
name='conv_f')
out_g = PF.convolution(x,
self.num_features,
kernel_shape,
pad=(pad_dim_0, pad_dim_1),
apply_w=mask_type,
name='conv_g')
if isinstance(payload, nn.Variable):
out_f += payload[:, :self.num_features, :, :]
out_g += payload[:, self.num_features:, :, :]
if self.conditional:
h_out_f = PF.affine(h, self.num_features, name='h_out_f')
h_out_f = h_out_f.reshape(
(h_out_f.shape[0], h_out_f.shape[1], 1, 1))
h_out_g = PF.affine(h, self.num_features, name='h_out_g')
h_out_g = h_out_g.reshape(
(h_out_g.shape[0], h_out_g.shape[1], 1, 1))
out = F.tanh(out_f + h_out_f) * F.sigmoid(out_g + h_out_g)
else:
out = F.tanh(out_f) * F.sigmoid(out_g)
if return_payload:
payload = PF.convolution(F.concatenate(out_f,
out_g,
axis=1),
2 * self.num_features, (1, 1),
name='conv_1x1')
payload = F.relu(payload)
return out, payload
else:
out = PF.convolution(x,
self.num_features,
kernel_shape,
stride=(1, 1),
pad=(pad_dim_0, pad_dim_1),
apply_w=mask_type)
out = F.relu(out)
return out
def network_LSTM(x, D, C, InputShape, HiddenSize, test=False):
# Input_2:x -> 687
# Delya_in:D -> 100
# Cell_in:C -> 100
# Concatenate -> 787
h = F.concatenate(D, x, axis=1)
# Affine -> 100
h1 = PF.affine(h, HiddenSize, name='Affine')
# InputGate -> 100
h2 = PF.affine(h, HiddenSize, name='InputGate')
# OutputGate -> 100
h3 = PF.affine(h, HiddenSize, name='OutputGate')
# ForgetGate -> 100
h4 = PF.affine(h, HiddenSize, name='ForgetGate')
# Sigmoid
h1 = F.sigmoid(h1)
# Sigmoid_2
h2 = F.sigmoid(h2)
# Sigmoid_3
h3 = F.sigmoid(h3)
# Sigmoid_4
h4 = F.sigmoid(h4)
# Mul2 -> 100
h1 = F.mul2(h1, h2)
# Mul2_3 -> 100
h4 = F.mul2(h4, C)
# Add2 -> 100
h1 = F.add2(h1, h4, True)
# Tanh
h5 = F.tanh(h1)
# Cell_out
h6 = F.identity(h1)
# Mul2_2 -> 100
h5 = F.mul2(h5, h3)
# Dropout
if not test:
h5 = F.dropout(h5)
# Output
h5 = F.identity(h5)
# Concatenate_2 -> 200
h5 = F.concatenate(h5, h6, axis=1)
return h5
def lstm_cell(x: nn.Variable, c: nn.Variable, h: nn.Variable) -> nn.Variable:
batch_size, units = c.shape
_hidden = PF.affine(F.concatenate(x, h, axis=1), 4*units)
a = F.tanh (_hidden[:, units*0: units*1])
input_gate = F.sigmoid(_hidden[:, units*1: units*2])
forgate_gate = F.sigmoid(_hidden[:, units*2: units*3])
output_gate = F.sigmoid(_hidden[:, units*3: units*4])
cell = input_gate * a + forgate_gate * c
hidden = output_gate * F.tanh(cell)
return cell, hidden
def lstm_cell(x, c, h):
batch_size, units = c.shape
_hidden = PF.affine(F.concatenate(x, h, axis=1), 4*units)
a = F.tanh (F.slice(_hidden, start=(0, units*0), stop=(batch_size, units*1)))
input_gate = F.sigmoid(F.slice(_hidden, start=(0, units*1), stop=(batch_size, units*2)))
forgate_gate = F.sigmoid(F.slice(_hidden, start=(0, units*2), stop=(batch_size, units*3)))
output_gate = F.sigmoid(F.slice(_hidden, start=(0, units*3), stop=(batch_size, units*4)))
cell = input_gate * a + forgate_gate * c
hidden = output_gate * F.tanh(cell)
return cell, hidden
def create_network(batchsize, imheight, imwidth, args, seen):
import gc
gc.collect()
nnabla_ext.cuda.clear_memory_cache()
anchors = args.num_anchors
classes = args.num_classes
yolo_x = nn.Variable((batchsize, 3, imheight, imwidth))
target = nn.Variable((batchsize, 50 * 5))
yolo_features = yolov2.yolov2(yolo_x, anchors, classes, test=False)
nB = yolo_features.shape[0]
nA = args.num_anchors
nC = args.num_classes
nH = yolo_features.shape[2]
nW = yolo_features.shape[3]
# Bouding box regression loss
# pred.shape = [nB, nA, 4, nH, nW]
output = F.reshape(yolo_features, (nB, nA, (5 + nC), nH, nW))
xy = F.sigmoid(output[:, :, :2, ...])
wh = output[:, :, 2:4, ...]
bbox_pred = F.concatenate(xy, wh, axis=2)
conf_pred = F.sigmoid(output[:, :, 4:5, ...])
cls_pred = output[:, :, 5:, ...]
region_loss_targets = RegionLossTargets(nC, args.anchors, seen,
args.coord_scale,
args.noobject_scale,
args.object_scale,
args.class_scale, args.thresh)
tcoord, mcoord, tconf, mconf, tcls, mcls = region_loss_targets(
bbox_pred, target)
for v in tcoord, mcoord, tconf, mconf, tcls, mcls:
v.need_grad = False
# Bounding box regression
bbox_loss = F.sum(F.squared_error(bbox_pred, tcoord) * mcoord)
# Conf (IoU) regression loss
conf_loss = F.sum(F.squared_error(conf_pred, tconf) * mconf)
# Class probability regression loss
cls_loss = F.sum(F.softmax_cross_entropy(cls_pred, tcls, axis=2) * mcls)
# Note:
# loss is devided by 2.0 due to the fact that the original darknet
# code doesn't multiply the derivative of square functions by 2.0
# in region_layer.c.
loss = (bbox_loss + conf_loss) / 2.0 + cls_loss
return yolo_x, target, loss, region_loss_targets
def propagate(h,
edges,
state_size=None,
w_initializer=None,
u_initializer1=None,
u_initializer2=None,
bias_initializer=None,
edge_initializers=None):
"""
Propagate vertex representations
Arguments:
h -- the input vertex representations (nnabla.Variable with shape (|V|, D))
edges -- the dictionary that represents the graph edge ({label, [in, out]})
state_size -- (optional) the size of hidden state (h.shape[1] is used if this argument is None)
w_initializer -- (optional)
u_initializer1 -- (optional)
u_initializer2 -- (optional)
bias_initializer -- (optional)
edge_initializers -- (optional)
Return value
- Return a variable with shape (|V|, D)
"""
if state_size is None:
state_size = h.shape[1]
h_size = h.shape[1]
with nn.parameter_scope("activate"):
a = activate(h,
edges,
state_size,
bias_initializer=bias_initializer,
edge_initializers=edge_initializers)
with nn.parameter_scope("W_zr"):
ws = PF.affine(a, (3, h_size), with_bias=False, w_init=w_initializer)
(z1, r1, h_hat1) = split(ws, axis=1)
with nn.parameter_scope("U_zr"):
us = PF.affine(h, (2, state_size),
with_bias=False,
w_init=u_initializer1)
(z2, r2) = split(us, axis=1)
z = F.sigmoid(F.add2(z1, z2))
r = F.sigmoid(F.add2(r1, r2))
with nn.parameter_scope("U"):
h_hat2 = PF.affine(F.mul2(r, h),
state_size,
with_bias=False,
w_init=u_initializer2)
h_hat = F.tanh(F.add2(h_hat1, h_hat2))
return F.add2(F.sub2(h, F.mul2(z, h)), F.mul2(z, h_hat))
def LSTM(inputs, units, return_sequences=False, name='lstm'):
'''
A long short-term memory layer
Args:
inputs (nnabla.Variable): A shape of [B, SentenceLength, EmbeddingSize].
units (int): Dimensionality of the output space.
return_sequences (bool): Whether to return the last output. in the output sequence, or the full sequence.
Returns:
nn.Variable: A shape [B, SentenceLength, units].
or
nn.Variable: A shape [B, units]
'''
hs = []
batch_size = inputs.shape[0]
sentence_length = inputs.shape[1]
c0 = nn.Variable.from_numpy_array(np.zeros((batch_size, units)))
h0 = nn.Variable.from_numpy_array(np.zeros((batch_size, units)))
inputs = F.split(inputs, axis=1)
cell = c0
hidden = h0
with nn.parameter_scope(name):
for x in inputs:
a = F.tanh(
PF.affine(x, units, with_bias=False, name='Wa') +
PF.affine(hidden, units, name='Ra'))
input_gate = F.sigmoid(
PF.affine(x, units, with_bias=False, name='Wi') +
PF.affine(hidden, units, name='Ri'))
forgate_gate = F.sigmoid(
PF.affine(x, units, with_bias=False, name='Wf') +
PF.affine(hidden, units, name='Rf'))
cell = input_gate * a + forgate_gate * cell
output_gate = F.sigmoid(
PF.affine(x, units, with_bias=False, name='Wo') +
PF.affine(hidden, units, name='Ro'))
hidden = output_gate * F.tanh(cell)
if return_sequences:
hidden = F.reshape(hidden, (batch_size, 1, units))
hs.append(hidden)
if return_sequences:
hs = F.concatenate(*hs, axis=1)
hs = F.reshape(hs, (batch_size, sentence_length, units))
return hs
else:
return hs[-1]
def SLE(f_large, f_small, scope_name):
with nn.parameter_scope(scope_name):
def sn_w(w): return PF.spectral_norm(w, dim=0)
ada_pool_size = f_small.shape[2] // 4
h = F.average_pooling(f_small, (ada_pool_size, ada_pool_size), stride=(
ada_pool_size, ada_pool_size))
h = PF.convolution(
h, f_large.shape[1], (4, 4), apply_w=sn_w, with_bias=False, name="conv1")
# Following the official implementation, this implementation uses swish instead of LeakyReLU here.
h = h * F.sigmoid(h)
h = PF.convolution(
h, f_large.shape[1], (1, 1), apply_w=sn_w, with_bias=False, name="conv2")
h = F.sigmoid(h)
h = f_large * h
return h
def lstm(x, h, c, w, b, with_bias):
hidden_size = h.shape[1]
xh = F.concatenate(*(x, h), axis=1)
w0, w1, w2, w3 = F.split(w, axis=0)
b0 = b1 = b2 = b3 = None
if with_bias:
b0, b1, b2, b3 = F.split(b, axis=0)
i_t = F.affine(xh, F.transpose(w0, (1, 0)), b0)
f_t = F.affine(xh, F.transpose(w1, (1, 0)), b1)
g_t = F.affine(xh, F.transpose(w2, (1, 0)), b2)
o_t = F.affine(xh, F.transpose(w3, (1, 0)), b3)
c_t = F.sigmoid(f_t) * c + F.sigmoid(i_t) * F.tanh(g_t)
h_t = F.sigmoid(o_t) * F.tanh(c_t)
return h_t, c_t
def test_sink(seed):
rng = np.random.RandomState(seed)
v = nn.Variable((2, 3, 4), need_grad=True)
h0 = F.tanh(v)
h1 = F.sigmoid(v)
v.d = rng.randn(*v.shape).astype(np.float32)
# Create references
v.grad.zero()
h0.forward()
h1.forward()
h0.backward()
h1.backward() # v.grad is accumulated.
h0d = h0.d.copy()
h1d = h1.d.copy()
vg = v.g.copy()
# Reset values
h0.data.zero()
h1.data.zero()
v.grad.zero()
# Check if sink works
dummy = F.sink(h0, h1, one_input_grad=True)
dummy.forward()
dummy.backward()
assert np.all(h0d == h0.d)
assert np.all(h1d == h1.d)
assert np.all(vg == v.g)
def G(z, is_train=True):
z = F.reshape(z, [batch_size, z_dim, 1, 1])
with nn.parameter_scope('G'):
with nn.parameter_scope('deconv1'):
dc1 = PF.deconvolution(z, 256, (4, 4), with_bias=False)
dc1 = PF.batch_normalization(dc1, batch_stat=is_train)
dc1 = F.leaky_relu(dc1)
with nn.parameter_scope('deconv2'):
dc2 = PF.deconvolution(dc1,
128, (4, 4),
pad=(1, 1),
stride=(2, 2),
with_bias=False)
dc2 = PF.batch_normalization(dc2, batch_stat=is_train)
dc2 = F.leaky_relu(dc2)
with nn.parameter_scope('deconv3'):
dc3 = PF.deconvolution(dc2,
64, (4, 4),
pad=(1, 1),
stride=(2, 2),
with_bias=False)
dc3 = PF.batch_normalization(dc3, batch_stat=is_train)
dc3 = F.leaky_relu(dc3)
with nn.parameter_scope('deconv4'):
dc4 = PF.deconvolution(dc3,
32, (4, 4),
pad=(3, 3),
stride=(2, 2),
with_bias=False)
dc4 = PF.batch_normalization(dc4, batch_stat=is_train)
dc4 = F.leaky_relu(dc4)
with nn.parameter_scope('output'):
output = PF.convolution(dc4, 1, (3, 3), pad=(1, 1))
output = F.sigmoid(output)
return output
def transformer(train=True, droput_ratio=0.1):
x = nn.Variable((batch_size, max_len))
t = nn.Variable((batch_size, 1))
mask = get_mask(x)
with nn.parameter_scope('embedding_layer'):
# h = time_distributed(PF.embed)(x, vocab_size, embedding_size) * mask
h = token_embedding(x, vocab_size, embedding_size)
h = position_encoding(h)
if train:
h = F.dropout(h, p=droput_ratio)
for i in range(hopping_num):
with nn.parameter_scope(f'encoder_hopping_{i}'):
h = residual_normalization_wrapper(multihead_self_attention)(
h,
head_num,
mask=mask,
train=train,
dropout_ratio=droput_ratio)
h = residual_normalization_wrapper(positionwise_feed_forward)(
h, train=train, dropout_ratio=droput_ratio)
with nn.parameter_scope('output_layer'):
y = F.sigmoid(PF.affine(h[:, 0, :], 1))
accuracy = F.mean(F.equal(F.round(y), t))
loss = F.mean(F.binary_cross_entropy(y, t))
return x, y, t, accuracy, loss
def backward_impl(self, inputs, outputs, prop_down, accum):
# inputs: [inputs_fwd_graph] + [inputs_bwd_graph] or
# [inputs_fwd_graph] + [outputs_fwd_graph] + [inputs_bwd_graph]
# Inputs
x0 = inputs[0].data
y0 = inputs[1].data
dy = inputs[2].data
# Outputs
dx0 = outputs[0].data
# Grads of inputs
g_x0 = inputs[0].grad
g_y0 = inputs[1].grad
g_dy = inputs[2].grad
# Grads of outputs
g_dx0 = outputs[0].grad
if prop_down[0]:
s = F.sigmoid(x0)
#y_dev = dx0 / dy
y_dev = y0 + s * (1.0 - y0)
g_x0_ = g_dx0 * dy * (y_dev + s * (1.0 - s) *
(1.0 - y0) - s * y_dev)
if accum[0]:
g_x0 += g_x0_
else:
g_x0.copy_from(g_x0_)
if prop_down[2]:
inp = nn.Variable(x0.shape).apply(data=x0,
grad=g_dy,
need_grad=True)
out = nn.Variable(dy.shape).apply(data=y0, grad=g_dx0)
self.forward_func.backward([inp], [out], accum=[accum[2]])
def backward_impl(self, inputs, outputs, prop_down, accum):
# inputs: [inputs_fwd_graph] + [inputs_bwd_graph] or
# [inputs_fwd_graph] + [outputs_fwd_graph] + [inputs_bwd_graph]
# Inputs
x0 = inputs[0].data
dy = inputs[1].data
# Outputs
dx0 = outputs[0].data
# Grads of inputs
g_x0 = inputs[0].grad
g_dy = inputs[1].grad
# Grads of outputs
g_dx0 = outputs[0].grad
# w.r.t. x0
sigmoid = F.sigmoid(x0)
if prop_down[0]:
if accum[0]:
g_x0 += g_dx0 * dy * sigmoid * (sigmoid - 1.0)
else:
g_x0.copy_from(g_dx0 * dy * sigmoid * (sigmoid - 1.0))
# w.r.t. dy
if prop_down[1]:
inp = nn.Variable(x0.shape).apply(data=x0,
grad=g_dy,
need_grad=True)
out = nn.Variable(dy.shape).apply(grad=g_dx0)
self.forward_func.backward([inp], [out], accum=[accum[1]])
def debug(self, debugger, images, dets, output, scale=1):
detection = dets.copy()
detection[:, :, :4] *= self.opt.down_ratio
hm = output[0]
hm = F.sigmoid(hm)
if self.opt.channel_last:
hm = F.transpose(hm, (0, 3, 1, 2))
for i in range(1):
if self.opt.mixed_precision:
# Removing pad from input image for drawing
img = images[i][:, :, :3]
if not self.opt.channel_last:
img = images[i].transpose(1, 2, 0)
img = ((img * self.std + self.mean) * 255).astype(np.uint8)
pred = debugger.gen_colormap(hm[i].d)
debugger.add_blend_img(img, pred, 'pred_hm_{:.1f}'.format(scale))
debugger.add_img(img, img_id='out_pred_{:.1f}'.format(scale))
for k in range(len(dets[i])):
if detection[i, k, 4] > self.opt.vis_thresh:
debugger.add_coco_bbox(detection[i, k, :4],
detection[i, k, -1],
detection[i, k, 4],
img_id='out_pred_{:.1f}'.format(scale))
for j in range(hm[i].shape[0]):
hmap = hm[i][j].d
hmap *= 255
hmap = hmap.astype('uint8')
print("max at channel {}:".format(j), np.max(hmap))
hmap = cv2.applyColorMap(hmap, cv2.COLORMAP_JET)
debugger.add_img(
hmap, img_id='heatmap_{}_{:.1f}'.format(j, scale))
def gru(x, h, w, b, with_bias):
hidden_size = h.shape[1]
xh = F.concatenate(*(x, h), axis=1)
w0, w1, w2 = F.split(w, axis=0)
b0 = b1 = b2 = b3 = None
if with_bias:
b0, b1, b2, b3 = F.split(b, axis=0)
r_t = F.sigmoid(F.affine(xh, F.transpose(w0, (1, 0)), b0))
z_t = F.sigmoid(F.affine(xh, F.transpose(w1, (1, 0)), b1))
w2_0 = w2[:, :w2.shape[1]-hidden_size]
w2_1 = w2[:, w2.shape[1]-hidden_size:]
n_t = F.tanh(F.affine(x, F.transpose(w2_0, (1, 0)), b2) +
r_t*F.affine(h, F.transpose(w2_1, (1, 0)), b3))
h_t = (1-z_t)*n_t + z_t*h
return h_t
def test_FLOPsEstimator():
x = nn.Variable((1, 3, 12, 12))
y = PF.depthwise_convolution(x, kernel=(5, 5), with_bias=True)
t = PF.fused_batch_normalization(y)
z = F.relu6(F.sigmoid(PF.affine(t, (3, 3), base_axis=2) + 3))
z = F.global_average_pooling(z)
est = FLOPsEstimator()
assert est.predict(z) == 17644
def model(img, sf):
"""
Define JSInet model
"""
with nn.parameter_scope('Network'):
with nn.parameter_scope('local_contrast_enhancement'):
## ================= Local Contrast Enhancement Subnet ============================ ##
ch = 64
b = guided_filter(img, 5, 0.01)
n1 = conv_2d(b, ch, kernel=(3, 3), name='conv/0')
for i in range(4):
n1 = res_block(n1, ch, 'res_block/%d' % i)
n1 = F.relu(n1, inplace=True)
local_filter_2d = conv_2d(
n1, (9**2) * (sf**2), kernel=(3, 3),
name='conv_k') # [B, H, W, (9x9)*(sfxsf)]
# dynamic 2D upsampling with 2D local filters
pred_C = dyn_2d_up_operation(b, local_filter_2d, (9, 9), sf)
# local contrast mask
pred_C = 2 * F.sigmoid(pred_C)
## ================= Detail Restoration Subnet ============================ ##
ch = 64
d = F.div2(img, b + 1e-15)
with nn.parameter_scope('detail_restoration'):
n3 = conv_2d(d, ch, kernel=(3, 3), name='conv/0')
for i in range(4):
n3 = res_block(n3, ch, 'res_block/%d' % i)
if i == 0:
d_feature = n3
n3 = F.relu(n3, inplace=True)
# separable 1D filters
dr_k_h = conv_2d(n3, 41 * sf**2, kernel=(3, 3), name='conv_k_h')
dr_k_v = conv_2d(n3, 41 * sf**2, kernel=(3, 3), name='conv_k_v')
# dynamic separable upsampling with with separable 1D local filters
pred_D = dyn_sep_up_operation(d, dr_k_v, dr_k_h, 41, sf)
## ================= Image Reconstruction Subnet ============================ ##
with nn.parameter_scope('image_reconstruction'):
n4 = conv_2d(img, ch, kernel=(3, 3), name='conv/0')
for i in range(4):
if i == 1:
n4 = F.concatenate(n4, d_feature, axis=3)
n4 = res_block_concat(n4, ch, 'res_block/%d' % i)
else:
n4 = res_block(n4, ch, 'res_block/%d' % i)
n4 = F.relu(n4, inplace=True)
n4 = F.relu(conv_2d(n4, ch * sf * sf, kernel=(3, 3),
name='conv/1'),
inplace=True)
# (1,100,170,1024) -> (1,100,170,4,4,64) -> (1,100,4,170,4,64)
# pixel shuffle
n4 = depth_to_space(n4, sf)
pred_I = conv_2d(n4, 3, kernel=(3, 3), name='conv/2')
pred = F.add2(pred_I, pred_D, inplace=True) * pred_C
jsinet = namedtuple('jsinet', ['pred'])
return jsinet(pred)
def gru(x, h, parameters_dict):
xh = F.concatenate(*(x, h), axis=1)
w0_nn = parameters_dict.get('w0_nn', None)
w1_nn = parameters_dict.get('w1_nn', None)
w2_0_nn = parameters_dict.get('w2_0_nn', None)
w2_1_nn = parameters_dict.get('w2_1_nn', None)
b0 = parameters_dict.get('b0', None)
b1 = parameters_dict.get('b1', None)
b2 = parameters_dict.get('b2', None)
b3 = parameters_dict.get('b3', None)
r_t = F.sigmoid(F.affine(xh, w0_nn, b0))
z_t = F.sigmoid(F.affine(xh, w1_nn, b1))
n_t = F.tanh(
F.affine(x, w2_0_nn, b2) + r_t * F.affine(h, w2_1_nn, b3))
h_t = (1 - z_t) * n_t + z_t * h
return h_t
def __call__(self, x):
outs = {}
feats = {}
with ps("discriminator/patch_gan"):
# Create all scale inputs first.
inputs = [x]
for i in range(self.n_scales - 1):
inputs.append(
F.average_pooling(inputs[-1], (3, 3), (2, 2),
pad=(1, 1),
including_pad=False))
for i in range(self.n_scales):
# Get input in reverse order of its scale (from coarse to fine) to preserve discriminator indexes.
h = inputs.pop()
d_name = "d_{}".format(i)
feats[d_name] = {}
with ps(d_name):
# first layer
fdim = self.base_ndf
with ps("layer_0"):
h = self.pad_conv(h, fdim, stride=(2, 2))
h = F.leaky_relu(h, alpha=0.2, inplace=True)
feats[d_name]["layer_0"] = h
# middle
for j in range(1, self.n_layers - 1):
fdim = min(fdim * 2, 512)
layer_name = "layer_{}".format(j)
with ps(layer_name):
h = self.pad_conv(h, fdim, stride=(2, 2))
h = self.instance_norm_lrelu(h)
feats[d_name][layer_name] = h
# last 2 layers
fdim = min(fdim * 2, 512)
layer_name = "layer_{}".format(self.n_layers - 1)
with ps(layer_name):
h = self.pad_conv(h, fdim, stride=(1, 1))
h = self.instance_norm_lrelu(h)
feats[d_name][layer_name] = h
with nn.parameter_scope("last_layer"):
h = self.pad_conv(h, 1, stride=(1, 1))
if self.use_sigmoid:
h = F.sigmoid(h)
outs[d_name] = h
return outs, feats
def backward_impl(self, inputs, outputs, prop_down, accum):
# inputs: [inputs_fwd_graph] + [inputs_bwd_graph] or
# [inputs_fwd_graph] + [outputs_fwd_graph] + [inputs_bwd_graph]
# Inputs
x0 = inputs[0].data # logits
t0 = inputs[1].data # labels
dz = inputs[2].data # grad_input
# Outputs
dx0 = outputs[0].data
dt0 = outputs[1].data
# Grads of inputs
g_x0 = inputs[0].grad
g_t0 = inputs[1].grad
g_dz = inputs[2].grad
# Grads of outputs
g_dx0 = outputs[0].grad
g_dt0 = outputs[1].grad
# Computation
## w.r.t. x0
if prop_down[0]:
sigmoid = F.sigmoid(x0)
g_x0_ = g_dx0 * dz * sigmoid * (1.0 - sigmoid)
if accum[0]:
g_x0 += g_x0_
else:
g_x0.copy_from(g_x0_)
## w.r.t. t0 is not required
## w.r.t. dz
if prop_down[2]:
# Instable implementation since using `/ dz`
# g_dz_ = g_dx0 * dx0 / dz
# x0 and t0 are of the same shape, no need to call F.sigmoid_cross_entropy
g_dz_ = g_dx0 * (F.sigmoid(x0) - t0)
if accum[2]:
g_dz += g_dz_
else:
g_dz.copy_from(g_dz_)
def network(x, maxh=16, depth=8):
with nn.parameter_scope("net"):
# (1, 28, 28) --> (32, 16, 16)
with nn.parameter_scope("convIn"):
out = F.tanh(PF.convolution(x, maxh, (1, 1), with_bias=False))
for i in range(depth):
with nn.parameter_scope("conv" + str(i)):
out = F.tanh(PF.convolution(out, maxh, (1, 1),
with_bias=False))
with nn.parameter_scope("convOut"):
out = F.sigmoid(PF.convolution(out, 3, (1, 1), with_bias=False))
return out
def log_sigmoid_backward(inputs):
"""
Args:
inputs (list of nn.Variable): Incomming grads/inputs to/of the forward function.
kwargs (dict of arguments): Dictionary of the corresponding function arguments.
Return:
list of Variable: Return the gradients wrt inputs of the corresponding function.
"""
dy = inputs[0]
x0 = inputs[1]
dx0 = dy * (1 - F.sigmoid(x0))
return dx0
def __init__(self,
batch_size=32,
learning_rate=1e-4,
max_iter=5086,
total_epochs=20,
monitor_path=None,
val_weight=None,
model_load_path=None):
"""
Construct all the necessary attributes for the attribute classifier.
Args:
batch_size (int): number of samples contained in each generated batch
learning_rate (float) : learning rate
max_iter (int) : maximum iterations for an epoch
total_epochs (int) : total epochs to train the model
val_weight : sample weights
monitor_path (str) : model parameter to be saved
model_load_path (str) : load the model
"""
self.batch_size = batch_size
# Resnet 50
# training graph
model = ResNet50()
self.input_image = nn.Variable((self.batch_size, ) + model.input_shape)
self.label = nn.Variable([self.batch_size, 1])
# fine tuning
pool = model(self.input_image, training=True, use_up_to='pool')
self.clf = clf_resnet50(pool)
self.clf.persistent = True
# loss
self.loss = F.mean(F.sigmoid_cross_entropy(self.clf, self.label))
# hyper parameters
self.solver = S.Adam(learning_rate)
self.solver.set_parameters(nn.get_parameters())
# validation graph
self.x_v = nn.Variable((self.batch_size, ) + model.input_shape)
pool_v = model(self.x_v, training=False, use_up_to='pool')
self.v_clf = clf_resnet50(pool_v, train=False)
self.v_clf_out = F.sigmoid(self.v_clf)
self.print_freq = 100
self.validation_weight = val_weight
# val params
self.acc = 0.0
self.total_epochs = total_epochs
self.max_iter = max_iter
self.monitor_path = monitor_path
if model_load_path is not None:
_ = nn.load_parameters(model_load_path)
def call(self, x, spk_emb, dilation):
dim = x.shape[1]
with nn.parameter_scope('shortcut'):
s = wn_conv(x, dim, (1, ))
with nn.parameter_scope('block'):
b = F.pad(x, (0, 0, dilation, dilation), 'reflect')
b = wn_conv(b,
2 * dim, (3, ),
dilation=(dilation, ),
name='conv_1')
if spk_emb is not None:
b = b + wn_conv(spk_emb, 2 * dim, (1, ), name="spk_emb")
b = F.tanh(b[:, :dim, ...]) * F.sigmoid(b[:, dim:, ...])
b = wn_conv(b, dim, (1, ), dilation=(dilation, ), name='conv_2')
return s + b
def sga(x):
return x * F.sigmoid(x)