MCPcopy Create free account
hub / github.com/pytorch/pytorch / UnrolledCell

Class UnrolledCell

caffe2/python/rnn_cell.py:1586–1662  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1584
1585
1586class UnrolledCell(RNNCell):
1587 def __init__(self, cell, T):
1588 self.T = T
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 #

Callers 1

_LSTMFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…