|
|
import mxnet as mx |
|
|
from nowcasting.ops import * |
|
|
from nowcasting.operators.common import identity, grid_generator, group_add |
|
|
from nowcasting.operators.base_rnn import MyBaseRNNCell |
|
|
import numpy as np |
|
|
|
|
|
|
|
|
|
|
|
class BaseConvRNN(MyBaseRNNCell): |
|
|
def __init__(self, num_filter, b_h_w, |
|
|
h2h_kernel=(3, 3), h2h_dilate=(1, 1), |
|
|
i2h_kernel=(3, 3), i2h_stride=(1, 1), |
|
|
i2h_pad=(1, 1), i2h_dilate=(1, 1), |
|
|
act_type="tanh", prefix="ConvRNN", params=None): |
|
|
super(BaseConvRNN, self).__init__(prefix=prefix + "_", params=params) |
|
|
self._num_filter = num_filter |
|
|
self._h2h_kernel = h2h_kernel |
|
|
assert (self._h2h_kernel[0] % 2 == 1) and (self._h2h_kernel[1] % 2 == 1), \ |
|
|
"Only support odd number, get h2h_kernel= %s" % str(h2h_kernel) |
|
|
self._h2h_pad = (h2h_dilate[0] * (h2h_kernel[0] - 1) // 2, |
|
|
h2h_dilate[1] * (h2h_kernel[1] - 1) // 2) |
|
|
self._h2h_dilate = h2h_dilate |
|
|
self._i2h_kernel = i2h_kernel |
|
|
self._i2h_stride = i2h_stride |
|
|
self._i2h_pad = i2h_pad |
|
|
self._i2h_dilate = i2h_dilate |
|
|
self._act_type = act_type |
|
|
assert len(b_h_w) == 3 |
|
|
i2h_dilate_ksize_h = 1 + (self._i2h_kernel[0] - 1) * self._i2h_dilate[0] |
|
|
i2h_dilate_ksize_w = 1 + (self._i2h_kernel[1] - 1) * self._i2h_dilate[1] |
|
|
self._batch_size, self._height, self._width = b_h_w |
|
|
self._state_height = (self._height + 2 * self._i2h_pad[0] - i2h_dilate_ksize_h)\ |
|
|
// self._i2h_stride[0] + 1 |
|
|
self._state_width = (self._width + 2 * self._i2h_pad[1] - i2h_dilate_ksize_w) \ |
|
|
// self._i2h_stride[1] + 1 |
|
|
print(self._prefix, self._state_height, self._state_width) |
|
|
self._curr_states = None |
|
|
self._counter = 0 |
|
|
|
|
|
|
|
|
class ConvRNN(BaseConvRNN): |
|
|
def __init__(self, num_filter, b_h_w, |
|
|
h2h_kernel=(3, 3), h2h_dilate=(1, 1), |
|
|
i2h_kernel=(3, 3), i2h_stride=(1, 1), |
|
|
i2h_pad=(1, 1), i2h_dilate=(1, 1), |
|
|
act_type="leaky", |
|
|
layer_norm=False, |
|
|
prefix="ConvRNN", |
|
|
params=None): |
|
|
super(ConvRNN, self).__init__(num_filter=num_filter, |
|
|
b_h_w=b_h_w, |
|
|
h2h_kernel=h2h_kernel, |
|
|
h2h_dilate=h2h_dilate, |
|
|
i2h_kernel=i2h_kernel, |
|
|
i2h_pad=i2h_pad, |
|
|
i2h_dilate=i2h_dilate, |
|
|
act_type=act_type, |
|
|
prefix=prefix, |
|
|
params=params) |
|
|
self._layer_norm = layer_norm |
|
|
self.i2h_weight = self.params.get('i2h_weight') |
|
|
self.i2h_bias = self.params.get('i2h_bias') |
|
|
self.h2h_weight = self.params.get('h2h_weight') |
|
|
self.h2h_bias = self.params.get('h2h_bias', init=mx.init.Normal()) |
|
|
|
|
|
@property |
|
|
def state_info(self): |
|
|
return [{'shape': (self._batch_size, self._num_filter, |
|
|
self._state_height, self._state_width), |
|
|
'__layout__': "NCHW"}] |
|
|
|
|
|
def __call__(self, inputs, states=None, is_initial=False, ret_mid=False): |
|
|
name = '%s_t%d' % (self._prefix, self._counter) |
|
|
self._counter += 1 |
|
|
states = self.begin_state()[0] if is_initial else states[0] |
|
|
assert states is not None |
|
|
if inputs is not None: |
|
|
i2h = mx.sym.Convolution(data=inputs, |
|
|
weight=self.i2h_weight, |
|
|
bias=self.i2h_bias, |
|
|
kernel=self._i2h_kernel, |
|
|
stride=self._i2h_stride, |
|
|
dilate=self._i2h_dilate, |
|
|
pad=self._i2h_pad, |
|
|
num_filter=self._num_filter, |
|
|
name="%s_i2h" % name) |
|
|
else: |
|
|
i2h = None |
|
|
h2h = mx.sym.Convolution(data=states, |
|
|
weight=self.h2h_weight, |
|
|
bias=self.h2h_bias, |
|
|
kernel=self._h2h_kernel, |
|
|
stride=(1, 1), |
|
|
dilate=self._h2h_dilate, |
|
|
pad=self._h2h_pad, |
|
|
num_filter=self._num_filter, |
|
|
name="%s_h2h" % name) |
|
|
if i2h is not None: |
|
|
if self._layer_norm: |
|
|
next_h = activation(layer_normalization(i2h + h2h, |
|
|
num_filters=self._num_filter, |
|
|
name=self._prefix + "ln"), |
|
|
act_type=self._act_type, name=name + "_state") |
|
|
else: |
|
|
next_h = activation(i2h + h2h, |
|
|
act_type=self._act_type, name=name + "_state") |
|
|
else: |
|
|
if self._layer_norm: |
|
|
next_h = activation(layer_normalization(h2h, |
|
|
num_filters=self._num_filter, |
|
|
name=self._prefix + "ln"), |
|
|
act_type=self._act_type, name=name + "_state") |
|
|
else: |
|
|
next_h = activation(h2h, act_type=self._act_type, name=name + "_state") |
|
|
|
|
|
self._curr_states = [next_h] |
|
|
if not ret_mid: |
|
|
return next_h, [next_h] |
|
|
else: |
|
|
return next_h, [next_h], [i2h, h2h] |
|
|
|
|
|
|
|
|
class ConvGRU(BaseConvRNN): |
|
|
def __init__(self, num_filter, b_h_w, zoneout=0.0, |
|
|
h2h_kernel=(3, 3), h2h_dilate=(1, 1), |
|
|
i2h_kernel=(3, 3), i2h_stride=(1, 1), i2h_pad=(1, 1), i2h_dilate=(1, 1), |
|
|
i2h_adj=(0, 0), no_i2h_bias=False, use_deconv=False, |
|
|
act_type="leaky", prefix="ConvGRU", lr_mult=1.0): |
|
|
"""Initializing a ConvGRU/DeconvGRU |
|
|
|
|
|
r_t = \sigma(W_r \ast x_t + R_r \ast h_{t-1} + b_{W_r} + b_{R_r}) |
|
|
u_t = \sigma(W_u \ast x_t + R_u \ast h_{t-1} + b_{W_u} + b_{R_u}) |
|
|
h^\prime_t = tanh(W_h \ast x_t + r_t \circ (R_h \ast h_{t-1} + b_{R_h}) + b_{W_h}) |
|
|
h_t = (1 - u_t) \circ h^\prime_t + u_t \circ h_{t-1} |
|
|
|
|
|
Parameters: (reset_gate, update_gate, new_mem) |
|
|
W_{i2h} = [W_r, W_u, W_h] |
|
|
b_{i2h} = [b_{W_r}, b_{W_u}, b_{W_h}] |
|
|
W_{h2h} = [R_r, R_u, R_h] |
|
|
b_{h2h} = [b_{R_r}, b_{R_u}, b_{R_h}] |
|
|
|
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
num_hidden : int |
|
|
hidden_act_type : str |
|
|
name : str |
|
|
""" |
|
|
super(ConvGRU, self).__init__(num_filter=num_filter, |
|
|
b_h_w=b_h_w, |
|
|
h2h_kernel=h2h_kernel, |
|
|
h2h_dilate=h2h_dilate, |
|
|
i2h_kernel=i2h_kernel, |
|
|
i2h_pad=i2h_pad, |
|
|
i2h_stride=i2h_stride, |
|
|
i2h_dilate=i2h_dilate, |
|
|
act_type=act_type, |
|
|
prefix=prefix) |
|
|
self._no_i2h_bias = no_i2h_bias |
|
|
self._i2h_adj = i2h_adj |
|
|
self._use_deconv = use_deconv |
|
|
if self._no_i2h_bias: |
|
|
assert use_deconv |
|
|
self._zoneout = zoneout |
|
|
self.i2h_weight = self.params.get("i2h_weight", lr_mult=lr_mult) |
|
|
self.i2h_bias = self.params.get("i2h_bias", lr_mult=lr_mult) |
|
|
self.h2h_weight = self.params.get("h2h_weight", lr_mult=lr_mult) |
|
|
self.h2h_bias = self.params.get("h2h_bias", lr_mult=lr_mult) |
|
|
|
|
|
@property |
|
|
def state_postfix(self): |
|
|
return ['h'] |
|
|
|
|
|
@property |
|
|
def state_info(self): |
|
|
return [{'shape': (self._batch_size, self._num_filter, |
|
|
self._state_height, self._state_width), |
|
|
'__layout__': "NCHW"}] |
|
|
|
|
|
def __call__(self, inputs, states=None, is_initial=False, ret_mid=False): |
|
|
name = '%s_t%d' % (self._prefix, self._counter) |
|
|
self._counter += 1 |
|
|
if is_initial: |
|
|
states = self.begin_state()[0] |
|
|
else: |
|
|
states = states[0] |
|
|
assert states is not None |
|
|
if inputs is not None: |
|
|
if self._use_deconv: |
|
|
if self._no_i2h_bias: |
|
|
i2h = mx.sym.Deconvolution(data=inputs, |
|
|
weight=self.i2h_weight, |
|
|
kernel=self._i2h_kernel, |
|
|
stride=self._i2h_stride, |
|
|
pad=self._i2h_pad, |
|
|
adj=self._i2h_adj, |
|
|
no_bias=True, |
|
|
num_filter=self._num_filter * 3, |
|
|
name="%s_i2h" % name) |
|
|
else: |
|
|
i2h = mx.sym.Deconvolution(data=inputs, |
|
|
weight=self.i2h_weight, |
|
|
bias=self.i2h_bias, |
|
|
kernel=self._i2h_kernel, |
|
|
stride=self._i2h_stride, |
|
|
pad=self._i2h_pad, |
|
|
adj=self._i2h_adj, |
|
|
num_filter=self._num_filter * 3, |
|
|
name="%s_i2h" % name) |
|
|
else: |
|
|
i2h = mx.sym.Convolution(data=inputs, |
|
|
weight=self.i2h_weight, |
|
|
bias=self.i2h_bias, |
|
|
kernel=self._i2h_kernel, |
|
|
stride=self._i2h_stride, |
|
|
dilate=self._i2h_dilate, |
|
|
pad=self._i2h_pad, |
|
|
num_filter=self._num_filter * 3, |
|
|
name="%s_i2h" % name) |
|
|
i2h_slice = mx.sym.SliceChannel(i2h, num_outputs=3, axis=1) |
|
|
else: |
|
|
i2h_slice = None |
|
|
prev_h = states |
|
|
print("h2h_dilate=", self._h2h_dilate) |
|
|
h2h = mx.sym.Convolution(data=prev_h, |
|
|
weight=self.h2h_weight, |
|
|
bias=self.h2h_bias, |
|
|
no_bias=False, |
|
|
kernel=self._h2h_kernel, |
|
|
stride=(1, 1), |
|
|
dilate=self._h2h_dilate, |
|
|
pad=self._h2h_pad, |
|
|
num_filter=self._num_filter * 3, |
|
|
name="%s_h2h" % name) |
|
|
h2h_slice = mx.sym.SliceChannel(h2h, num_outputs=3, axis=1) |
|
|
if i2h_slice is not None: |
|
|
reset_gate = mx.sym.Activation(i2h_slice[0] + h2h_slice[0], act_type="sigmoid", |
|
|
name=name + "_r") |
|
|
update_gate = mx.sym.Activation(i2h_slice[1] + h2h_slice[1], act_type="sigmoid", |
|
|
name=name + "_u") |
|
|
new_mem = activation(i2h_slice[2] + reset_gate * h2h_slice[2], |
|
|
act_type=self._act_type, |
|
|
name=name + "_h") |
|
|
else: |
|
|
reset_gate = mx.sym.Activation(h2h_slice[0], act_type="sigmoid", |
|
|
name=name + "_r") |
|
|
update_gate = mx.sym.Activation(h2h_slice[1], act_type="sigmoid", |
|
|
name=name + "_u") |
|
|
new_mem = activation(reset_gate * h2h_slice[2], |
|
|
act_type=self._act_type, |
|
|
name=name + "_h") |
|
|
next_h = update_gate * prev_h + (1 - update_gate) * new_mem |
|
|
if self._zoneout > 0.0: |
|
|
mask = mx.sym.Dropout(mx.sym.ones_like(prev_h), p=self._zoneout) |
|
|
next_h = mx.sym.where(mask, next_h, prev_h) |
|
|
self._curr_states = [next_h] |
|
|
if not ret_mid: |
|
|
return next_h, [next_h] |
|
|
else: |
|
|
return next_h, [next_h], [] |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
import numpy as np |
|
|
|
|
|
|
|
|
data = mx.sym.Variable('data') |
|
|
data = mx.sym.SliceChannel(data, axis=0, num_outputs=11, squeeze_axis=True) |
|
|
conv_gru1 = ConvGRU(num_filter=100, b_h_w=(3, 40, 40), |
|
|
prefix="conv_gru1") |
|
|
out, states = conv_gru1(inputs=data[0], is_initial=True) |
|
|
for i in range(1, 11): |
|
|
out, states = conv_gru1(inputs=data[i], states=states) |
|
|
conv_gru_forward_backward_time =\ |
|
|
mx.test_utils.check_speed(out, |
|
|
location={'data': np.random.normal(size=(11, 3, 128, 40, 40))}, |
|
|
N=2) |
|
|
net = mx.mod.Module(out, data_names=['data',], label_names=None, context=mx.gpu()) |
|
|
net.bind(data_shapes=[('data', (11, 3, 128, 40, 40))], |
|
|
grad_req='add') |
|
|
net.init_params() |
|
|
net.forward(mx.io.DataBatch(data=[mx.random.normal(shape=(11, 3, 128, 40, 40))], label=None), is_train=False) |
|
|
print(net.get_outputs()[0].asnumpy()) |
|
|
|
|
|
|
|
|
data = mx.sym.Variable('data') |
|
|
data = mx.sym.SliceChannel(data, axis=0, num_outputs=11, squeeze_axis=True) |
|
|
conv_rnn1 = ConvRNN(num_filter=100, b_h_w=(3, 40, 40), |
|
|
prefix="conv_rnn1") |
|
|
out, states = conv_rnn1(inputs=data[0], is_initial=True) |
|
|
for i in range(1, 11): |
|
|
out, states = conv_rnn1(inputs=data[i], states=states) |
|
|
conv_rnn_forward_backward_time = \ |
|
|
mx.test_utils.check_speed(out, |
|
|
location={'data': np.random.normal(size=(11, 3, 128, 40, 40))}, |
|
|
N=2) |
|
|
net = mx.mod.Module(out, data_names=['data', ], label_names=None, context=mx.gpu()) |
|
|
net.bind(data_shapes=[('data', (11, 3, 128, 40, 40))], |
|
|
grad_req='add') |
|
|
net.init_params() |
|
|
net.forward(mx.io.DataBatch(data=[mx.random.normal(shape=(11, 3, 128, 40, 40))], label=None), |
|
|
is_train=False) |
|
|
print(net.get_outputs()[0].asnumpy()) |
|
|
|
|
|
print("ConvGRU Time:", conv_gru_forward_backward_time) |
|
|
print("ConvRNN Time:", conv_rnn_forward_backward_time) |
|
|
|
