def project_and_conv1d(d, noutput, r=5):
return [
layers.Fun("lambda x: x.max(2)[0]"),
flex.Conv1d(d, r),
flex.BatchNorm1d(),
nn.ReLU(),
flex.Conv1d(noutput, 1),
layers.Reorder("BDL", ocr_output)
]
def project_and_lstm(d, noutput, num_layers=1):
return [
layers.Fun("lambda x: x.sum(2)"), # BDHW -> BDW
layers.Reorder("BDL", "LBD"),
flex.LSTM(d, bidirectional=True, num_layers=num_layers),
layers.Reorder("LBD", "BDL"),
flex.Conv1d(noutput, 1),
layers.Reorder("BDL", ocr_output)
]
def make_lstm_transpose(noutput=noutput):
model = nn.Sequential(
layers.Input("BDHW", range=(0, 1), sizes=[None, 1, None, None]),
*combos.conv2d_block(50, 3, repeat=2),
*combos.conv2d_block(100, 3, repeat=2),
*combos.conv2d_block(150, 3, repeat=2),
*combos.conv2d_block(200, 3, repeat=2),
layers.Fun("lambda x: x.sum(2)"), # BDHW -> BDW
flex.ConvTranspose1d(800, 1, stride=2), # <-- undo too tight spacing
#flex.BatchNorm1d(), nn.ReLU(),
layers.Reorder("BDL", "LBD"),
flex.LSTM(100, bidirectional=True),
layers.Reorder("LBD", "BDL"),
flex.Conv1d(noutput, 1),
layers.Reorder("BDL", ocr_output))
flex.shape_inference(model, (1, 1, 128, 512))
return model
def make_lstm_keep(noutput=noutput):
model = nn.Sequential(
layers.Input("BDHW", range=(0, 1), sizes=[None, 1, None, None]),
layers.KeepSize(
mode="nearest",
dims=[3],
sub=nn.Sequential(
*combos.conv2d_block(50, 3, repeat=2),
*combos.conv2d_block(100, 3, repeat=2),
*combos.conv2d_block(150, 3, repeat=2),
layers.Fun("lambda x: x.sum(2)") # BDHW -> BDW
)),
flex.Conv1d(500, 5, padding=2),
flex.BatchNorm1d(),
nn.ReLU(),
layers.Reorder("BDL", "LBD"),
flex.LSTM(200, bidirectional=True),
layers.Reorder("LBD", "BDL"),
flex.Conv1d(noutput, 1),
layers.Reorder("BDL", ocr_output))
flex.shape_inference(model, (1, 1, 128, 512))
return model
def test_Fun():
mod = layers.Fun("lambda x: x[0]**2")
a = torch.ones((2, 3, 4))
b = mod(a)
assert (b == 1).all()
assert tuple(b.shape) == (3, 4)