|
|
import mxnet as mx |
|
|
from mxnet.rnn import BaseRNNCell |
|
|
from nowcasting.ops import activation |
|
|
from nowcasting.operators.common import group_add |
|
|
|
|
|
class MyBaseRNNCell(BaseRNNCell): |
|
|
def __init__(self, prefix="MyBaseRNNCell", params=None): |
|
|
super(MyBaseRNNCell, self).__init__(prefix=prefix, params=params) |
|
|
|
|
|
def __call__(self, inputs, states, is_initial=False, ret_mid=False): |
|
|
raise NotImplementedError() |
|
|
|
|
|
def reset(self): |
|
|
super(MyBaseRNNCell, self).reset() |
|
|
self._curr_states = None |
|
|
|
|
|
def get_current_states(self): |
|
|
return self._curr_states |
|
|
|
|
|
def unroll(self, length, inputs=None, begin_state=None, ret_mid=False, |
|
|
input_prefix='', layout='TC', merge_outputs=False): |
|
|
"""Unroll an RNN cell across time steps. |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
length : int |
|
|
number of steps to unroll |
|
|
inputs : Symbol, list of Symbol, or None |
|
|
if inputs is a single Symbol (usually the output |
|
|
of Embedding symbol), it should have shape |
|
|
(batch_size, length, ...) if layout == 'NTC', |
|
|
or (length, batch_size, ...) if layout == 'TNC'. |
|
|
|
|
|
If inputs is a list of symbols (usually output of |
|
|
previous unroll), they should all have shape |
|
|
(batch_size, ...). |
|
|
|
|
|
If inputs is None, Placeholder variables are |
|
|
automatically created. |
|
|
begin_state : nested list of Symbol |
|
|
input states. Created by begin_state() |
|
|
or output state of another cell. Created |
|
|
from begin_state() if None. |
|
|
input_prefix : str |
|
|
prefix for automatically created input |
|
|
placehodlers. |
|
|
layout : str |
|
|
layout of input symbol. Only used if inputs |
|
|
is a single Symbol. |
|
|
merge_outputs : bool |
|
|
if False, return outputs as a list of Symbols. |
|
|
If True, concatenate output across time steps |
|
|
and return a single symbol with shape |
|
|
(batch_size, length, ...) if layout == 'NTC', |
|
|
or (length, batch_size, ...) if layout == 'TNC'. |
|
|
|
|
|
Returns |
|
|
------- |
|
|
outputs : list of Symbol |
|
|
output symbols. |
|
|
states : Symbol or nested list of Symbol |
|
|
has the same structure as begin_state() |
|
|
mid_info : list of Symbol |
|
|
""" |
|
|
self.reset() |
|
|
assert layout == 'TNC' or layout == 'TC' |
|
|
if inputs is not None: |
|
|
if isinstance(inputs, mx.sym.Symbol): |
|
|
assert len(inputs.list_outputs()) == 1, \ |
|
|
"unroll doesn't allow grouped symbol as input. Please " \ |
|
|
"convert to list first or let unroll handle slicing" |
|
|
if 'N' in layout: |
|
|
inputs = mx.sym.SliceChannel(inputs, axis=0, num_outputs=length, |
|
|
squeeze_axis=1) |
|
|
else: |
|
|
inputs = mx.sym.SliceChannel(inputs, axis=0, num_outputs=length) |
|
|
else: |
|
|
assert len(inputs) == length |
|
|
else: |
|
|
inputs = [None] * length |
|
|
if begin_state is None: |
|
|
states = self.begin_state() |
|
|
else: |
|
|
states = begin_state |
|
|
outputs = [] |
|
|
mid_infos = [] |
|
|
for i in range(length): |
|
|
output, states, mid_info = self(inputs=inputs[i], states=states, |
|
|
is_initial=(i == 0 and (begin_state is None)), |
|
|
ret_mid=True) |
|
|
outputs.append(output) |
|
|
mid_infos.extend(mid_info) |
|
|
if merge_outputs: |
|
|
outputs = [mx.sym.expand_dims(i, axis=0) for i in outputs] |
|
|
outputs = mx.sym.Concat(*outputs, dim=0) |
|
|
if ret_mid: |
|
|
return outputs, states, mid_infos |
|
|
else: |
|
|
return outputs, states |
|
|
|
|
|
|
|
|
class BaseStackRNN(object): |
|
|
def __init__(self, base_rnn_class, stack_num=1, |
|
|
name="BaseStackRNN", residual_connection=True, |
|
|
**kwargs): |
|
|
self._base_rnn_class = base_rnn_class |
|
|
self._residual_connection = residual_connection |
|
|
self._name = name |
|
|
self._stack_num = stack_num |
|
|
self._prefix = name + "_" |
|
|
self._rnns = [base_rnn_class(prefix=self._name + "_%d" %i, **kwargs) for i in range(stack_num)] |
|
|
self._init_counter = 0 |
|
|
self._state_info = None |
|
|
|
|
|
def init_state_vars(self): |
|
|
"""Initial state variable for this cell. |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
|
|
|
Returns |
|
|
------- |
|
|
state_vars : nested list of Symbol |
|
|
starting states for first RNN step |
|
|
""" |
|
|
state_vars = [] |
|
|
for i, info in enumerate(self.state_info): |
|
|
state = mx.sym.var(name='%s_begin_state_%s' % (self._name, self.state_postfix[i]), **info) |
|
|
state_vars.append(state) |
|
|
return state_vars |
|
|
|
|
|
def concat_to_split(self, concat_states): |
|
|
assert len(concat_states) == len(self.state_info) |
|
|
split_states = [[] for i in range(self._stack_num)] |
|
|
for i in range(len(self.state_info)): |
|
|
channel_axis = self.state_info[i]['__layout__'].lower().find('c') |
|
|
ele = mx.sym.split(concat_states[i], num_outputs=self._stack_num, axis=channel_axis) |
|
|
for j in range(self._stack_num): |
|
|
split_states[j].append(ele[j]) |
|
|
return split_states |
|
|
|
|
|
def split_to_concat(self, split_states): |
|
|
|
|
|
concat_states = [] |
|
|
for i in range(len(self.state_info)): |
|
|
channel_axis = self.state_info[i]['__layout__'].lower().find('c') |
|
|
concat_states.append(mx.sym.concat(*[ele[i] for ele in split_states], |
|
|
dim=channel_axis)) |
|
|
return concat_states |
|
|
|
|
|
def check_concat(self, states): |
|
|
ret = not isinstance(states[0], list) |
|
|
return ret |
|
|
|
|
|
def to_concat(self, states): |
|
|
if not self.check_concat(states): |
|
|
states = self.split_to_concat(states) |
|
|
return states |
|
|
|
|
|
def to_split(self, states): |
|
|
if self.check_concat(states): |
|
|
states = self.concat_to_split(states) |
|
|
return states |
|
|
|
|
|
@property |
|
|
def state_postfix(self): |
|
|
return self._rnns[0].state_postfix |
|
|
|
|
|
@property |
|
|
def state_info(self): |
|
|
if self._state_info is None: |
|
|
info = [] |
|
|
for i in range(len(self._rnns[0].state_info)): |
|
|
ele = {} |
|
|
for rnn in self._rnns: |
|
|
if 'shape' not in ele: |
|
|
ele['shape'] = list(rnn.state_info[i]['shape']) |
|
|
else: |
|
|
channel_dim = rnn.state_info[i]['__layout__'].lower().find('c') |
|
|
ele['shape'][channel_dim] += rnn.state_info[i]['shape'][channel_dim] |
|
|
if '__layout__' not in ele: |
|
|
ele['__layout__'] = rnn.state_info[i]['__layout__'].upper() |
|
|
else: |
|
|
assert rnn.state_info[i]['__layout__'] == ele['__layout__'].upper() |
|
|
ele['shape'] = tuple(ele['shape']) |
|
|
info.append(ele) |
|
|
self._state_info = info |
|
|
return info |
|
|
else: |
|
|
return self._state_info |
|
|
|
|
|
def flatten_add_layout(self, states, blocked=False): |
|
|
""" |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
states : list of list or list |
|
|
|
|
|
Returns |
|
|
------- |
|
|
ret : list |
|
|
""" |
|
|
states = self.to_concat(states) |
|
|
assert self.check_concat(states) |
|
|
ret = [] |
|
|
for i, ele in enumerate(states): |
|
|
if blocked: |
|
|
ret.append(mx.sym.BlockGrad(ele, __layout__=self.state_info[i]['__layout__'])) |
|
|
else: |
|
|
ele._set_attr(__layout__=self.state_info[i]['__layout__']) |
|
|
ret.append(ele) |
|
|
return ret |
|
|
|
|
|
def reset(self): |
|
|
for i in range(len(self._rnns)): |
|
|
self._rnns[i].reset() |
|
|
|
|
|
def unroll(self, length, inputs=None, begin_states=None, ret_mid=False): |
|
|
if begin_states is None: |
|
|
begin_states = self.init_state_vars() |
|
|
begin_states = self.to_split(begin_states) |
|
|
assert len(begin_states) == self._stack_num, len(begin_states) |
|
|
for ele in begin_states: |
|
|
assert len(ele) == len(self.state_info) |
|
|
outputs = [] |
|
|
final_states = [] |
|
|
mid_infos = [] |
|
|
for i in range(len(self._rnns)): |
|
|
rnn_out_list, rnn_final_states, rnn_mid_infos =\ |
|
|
self._rnns[i].unroll(length=length, inputs=inputs, |
|
|
begin_state=begin_states[i], |
|
|
layout="TC", |
|
|
ret_mid=True) |
|
|
if self._residual_connection and i > 0: |
|
|
|
|
|
rnn_out_list = group_add(lhs=rnn_out_list, rhs=inputs) |
|
|
inputs = rnn_out_list |
|
|
outputs.append(rnn_out_list) |
|
|
final_states.append(rnn_final_states) |
|
|
mid_infos.append(rnn_mid_infos) |
|
|
if ret_mid: |
|
|
return outputs, final_states, mid_infos |
|
|
else: |
|
|
return outputs, final_states |
|
|
|
|
|
|
|
|
class MyGRU(MyBaseRNNCell): |
|
|
"""GRU cell. |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
num_hidden : int |
|
|
number of units in output symbol |
|
|
prefix : str, default 'rnn_' |
|
|
prefix for name of layers |
|
|
(and name of weight if params is None) |
|
|
params : RNNParams or None |
|
|
container for weight sharing between cells. |
|
|
created if None. |
|
|
""" |
|
|
def __init__(self, num_hidden, zoneout=0.0, act_type="tanh", prefix='gru_', params=None): |
|
|
super(MyGRU, self).__init__(prefix=prefix, params=params) |
|
|
self._num_hidden = num_hidden |
|
|
self._act_type = act_type |
|
|
self._zoneout = zoneout |
|
|
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') |
|
|
|
|
|
@property |
|
|
def state_info(self): |
|
|
"""shape(s) of states""" |
|
|
return [{'shape': (0, self._num_hidden), '__layout__': "NC"}] |
|
|
|
|
|
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: |
|
|
prev_h = self.begin_state()[0] |
|
|
else: |
|
|
prev_h = states[0] |
|
|
assert states is not None |
|
|
if inputs is not None: |
|
|
inputs = mx.sym.reshape(inputs, shape=(0, -1)) |
|
|
i2h = mx.sym.FullyConnected(data=inputs, |
|
|
num_hidden=self._num_hidden * 3, |
|
|
weight=self._i2h_weight, |
|
|
bias=self._i2h_bias, |
|
|
name="%s_i2h" %name) |
|
|
i2h_slice = mx.sym.SliceChannel(i2h, num_outputs=3, axis=1) |
|
|
else: |
|
|
i2h_slice = None |
|
|
h2h = mx.sym.FullyConnected(data=prev_h, |
|
|
num_hidden=self._num_hidden * 3, |
|
|
weight=self._h2h_weight, |
|
|
bias=self._h2h_bias, |
|
|
name="%s_h2h" %name) |
|
|
h2h_slice = mx.sym.SliceChannel(h2h, num_outputs=3, axis=1) |
|
|
if i2h_slice is not None: |
|
|
reset_gate = activation(i2h_slice[0] + h2h_slice[0], act_type="sigmoid", |
|
|
name=name + "_r") |
|
|
update_gate = 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 = activation(h2h_slice[0], act_type="sigmoid", |
|
|
name=name + "_r") |
|
|
update_gate = 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__': |
|
|
from nowcasting.operators.conv_rnn import ConvGRU |
|
|
brnn1 = BaseStackRNN(base_rnn_class=ConvGRU, stack_num=5, |
|
|
b_h_w=(4, 32, 32), num_filter=32) |
|
|
print(brnn1.state_info) |
|
|
inputs = mx.sym.var(name="inputs", shape=(8, 4, 16, 32, 32)) |
|
|
outputs, final_states, mid_infos = brnn1.unroll(length=8, inputs=inputs, ret_mid=True) |
|
|
print(len(outputs), len(outputs[0])) |
|
|
print(len(final_states), len(final_states[0])) |
