SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 3 Recurrent Acceleration Network (RAN)
for layer in range(self.num_layers):
    if layer == 0:
        a = input_transfer.reshape(Water2sea.slice_num * batch_size, Water2sea.slice_len, input_size)
        W = self.W_first_layer
    else:
        a = F.dropout(output_all Slice, self.dropout, self.training)
        if layer == 1:
            layer0_output = a
            W = self.W_other_layer(layer-1, :, :]
            hidden.slice_row = hidden.slice_row * 0
            hidden.slice_col = hidden.slice_col * 0
        B = self.B(layer, :]
    # start every for all slice
    output_all.slice_list = []
    for slice in range (Water2sea.slice_len):
        # gate generate
        gate = self.Linear(torch.cat([hidden.slice_row, hidden.slice_col, a(:, slice, :],
                    dim = -1), W, B, batch_size, slice, Water2sea.slice_num)
        # gate
        sigmoid_GATE, tanh_GATE = torch.split(gate, 4 * self-hidden_size, dim = -1)
        sigmoid_GATE = torch.sigmoid(sigmoid_GATE)
        tanh_GATE = torch.tanh(tanh_GATE)
        update_GATE_row, output_GATE_row, update_GATE_col, output_GATE_col = sigmod_GATE.chunk(4, dim = -1)
        input_GATE_row, input_GATE_col = tanh_GATE.chunk(2, dim = -1)
        # gate effect
        hidden.slice_row = torch.tanh(
            (1-update_GATE_row)*hidden.slice_row + update_GATE_row*input_GATE_row) *
            output_GATE_row
        hidden.slice_col = torch.tanh(
            (1-update_GATE_col)*hidden.slice_col + update_GATE_col*input_GATE_col) *
            output_GATE_col
        # output generate
        output Slice = torch.cat([hidden.slice_row, hidden.slice_col], dim = -1)
        # save output
        output_all Slice_list.append(output Slice)
        # save row hidden
        if slice >= Original Slice_len - 1:
            need_save_row_LOC = slice - Original Slice_len + 1
            hidden_row_all_list.append(
                hidden.slice_row[need_save_row_LOC*batch_size:(need_save_row_LOC+1)*batch_size, :])
        # save col hidden
        if slice >= Water2sea.slice_num - 1:
            hidden_col_all_list.append(
                hidden SCElice_col[(Water2sea.slice_num-1)*batch_size,:])
        # hidden transfer
        hidden SCElice_col = torch-roll(hidden SCElice_col, shifts=batch_size, dims = 0)
        if self.res_mode == 'layer_res' and layer >= 1: # layer-res
            output_all Slice = torch.stack(output_all SCElice_list, dim = 1) + layer0_output
    else:
        output_all Slice = torch.stack(output_all SCElice_list, dim = 1)
        hidden_row_all = torch.stack(hidden_row_all_list, dim = 1)
        hidden_col_all = torch.stack(hidden_col_all_list, dim = 1)
        hidden_row_all = hidden_row_all.reshape(batch_size, self.num_layers, Water2sea.slice_num,
                     hidden_row_all.shape[-1])
        hidden_col_all = hidden_col_allreshape(batch_size, self.num_layers, Original SCElice_len,
                     hidden_col_all.shape[-1])
    if flag == 1:
        return output_all Slice, hidden_col_all, hidden_row_all
    else:
        return output_all Slice, hidden_row_all, hidden_col_all