| import mxnet as mx |
| from nowcasting.config import cfg |
| from nowcasting.ops import reset_regs |
| from nowcasting.operators.common import grid_generator |
|
|
|
|
| class PredictionBaseFactory(object): |
| def __init__(self, batch_size, in_seq_len, out_seq_len, height, width, name="forecaster"): |
| self._out_typ = cfg.MODEL.OUT_TYPE |
| self._batch_size = batch_size |
| self._in_seq_len = in_seq_len |
| self._out_seq_len = out_seq_len |
| self._height = height |
| self._width = width |
| self._name = name |
| self._spatial_grid = grid_generator(batch_size=batch_size, height=height, width=width) |
| self.rnn_list = self._init_rnn() |
| self._reset_rnn() |
|
|
| def _pre_encode_frame(self, frame_data, seqlen): |
| ret = mx.sym.Concat(frame_data, |
| mx.sym.broadcast_to(mx.sym.expand_dims(self._spatial_grid, axis=0), |
| shape=(seqlen, self._batch_size, |
| 2, self._height, self._width)), |
| mx.sym.ones(shape=(seqlen, self._batch_size, 1, |
| self._height, self._width)), |
| num_args=3, dim=2) |
| return ret |
|
|
| def _init_rnn(self): |
| raise NotImplementedError |
|
|
| def _reset_rnn(self): |
| for rnn in self.rnn_list: |
| rnn.reset() |
|
|
| def reset_all(self): |
| reset_regs() |
| self._reset_rnn() |
|
|
|
|
| class RecursiveOneStepBaseFactory(PredictionBaseFactory): |
| def __init__(self, batch_size, in_seq_len, out_seq_len, height, width, use_ss=False, |
| name="forecaster"): |
| super(RecursiveOneStepBaseFactory, self).__init__(batch_size=batch_size, |
| in_seq_len=in_seq_len, |
| out_seq_len=out_seq_len, |
| height=height, |
| width=width, |
| name=name) |
| self._use_ss = False |
|
|
|
|
