| 1589 | self.cell = cell |
| 1590 | |
| 1591 | def apply_over_sequence( |
| 1592 | self, |
| 1593 | model, |
| 1594 | inputs, |
| 1595 | seq_lengths, |
| 1596 | initial_states, |
| 1597 | outputs_with_grads=None, |
| 1598 | ): |
| 1599 | inputs = self.cell.prepare_input(model, inputs) |
| 1600 | |
| 1601 | # Now they are blob references - outputs of splitting the input sequence |
| 1602 | split_inputs = model.net.Split( |
| 1603 | inputs, |
| 1604 | [str(inputs) + "_timestep_{}".format(i) |
| 1605 | for i in range(self.T)], |
| 1606 | axis=0) |
| 1607 | if self.T == 1: |
| 1608 | split_inputs = [split_inputs] |
| 1609 | |
| 1610 | states = initial_states |
| 1611 | all_states = [] |
| 1612 | for t in range(0, self.T): |
| 1613 | scope_name = "timestep_{}".format(t) |
| 1614 | # Parameters of all timesteps are shared |
| 1615 | with ParameterSharing({scope_name: ''}),\ |
| 1616 | scope.NameScope(scope_name): |
| 1617 | timestep = model.param_init_net.ConstantFill( |
| 1618 | [], "timestep", value=t, shape=[1], |
| 1619 | dtype=core.DataType.INT32, |
| 1620 | device_option=core.DeviceOption(caffe2_pb2.CPU)) |
| 1621 | states = self.cell._apply( |
| 1622 | model=model, |
| 1623 | input_t=split_inputs[t], |
| 1624 | seq_lengths=seq_lengths, |
| 1625 | states=states, |
| 1626 | timestep=timestep, |
| 1627 | ) |
| 1628 | all_states.append(states) |
| 1629 | |
| 1630 | all_states = zip(*all_states) |
| 1631 | all_states = [ |
| 1632 | model.net.Concat( |
| 1633 | list(full_output), |
| 1634 | [ |
| 1635 | str(full_output[0])[len("timestep_0/"):] + "_concat", |
| 1636 | str(full_output[0])[len("timestep_0/"):] + "_concat_info" |
| 1637 | |
| 1638 | ], |
| 1639 | axis=0)[0] |
| 1640 | for full_output in all_states |
| 1641 | ] |
| 1642 | # Interleave the state values similar to |
| 1643 | # |
| 1644 | # x = [1, 3, 5] |
| 1645 | # y = [2, 4, 6] |
| 1646 | # z = [val for pair in zip(x, y) for val in pair] |
| 1647 | # # z is [1, 2, 3, 4, 5, 6] |
| 1648 | # |