(args, queue, label_queue, input_shape)
| 65 | |
| 66 | |
| 67 | def create_model(args, queue, label_queue, input_shape): |
| 68 | model = model_helper.ModelHelper(name="LSTM_bench") |
| 69 | seq_lengths, target = \ |
| 70 | model.net.AddExternalInputs( |
| 71 | 'seq_lengths', |
| 72 | 'target', |
| 73 | ) |
| 74 | |
| 75 | input_blob = model.net.DequeueBlobs(queue, "input_data") |
| 76 | labels = model.net.DequeueBlobs(label_queue, "label") |
| 77 | |
| 78 | init_blobs = [] |
| 79 | if args.implementation in ["own", "static", "static_dag"]: |
| 80 | T = None |
| 81 | if "static" in args.implementation: |
| 82 | assert args.fixed_shape, \ |
| 83 | "Random input length is not static RNN compatible" |
| 84 | T = args.seq_length |
| 85 | print("Using static RNN of size {}".format(T)) |
| 86 | |
| 87 | for i in range(args.num_layers): |
| 88 | hidden_init, cell_init = model.net.AddExternalInputs( |
| 89 | "hidden_init_{}".format(i), |
| 90 | "cell_init_{}".format(i) |
| 91 | ) |
| 92 | init_blobs.extend([hidden_init, cell_init]) |
| 93 | |
| 94 | output, last_hidden, _, last_state = rnn_cell.LSTM( |
| 95 | model=model, |
| 96 | input_blob=input_blob, |
| 97 | seq_lengths=seq_lengths, |
| 98 | initial_states=init_blobs, |
| 99 | dim_in=args.input_dim, |
| 100 | dim_out=[args.hidden_dim] * args.num_layers, |
| 101 | scope="lstm1", |
| 102 | memory_optimization=args.memory_optimization, |
| 103 | forward_only=args.forward_only, |
| 104 | drop_states=True, |
| 105 | return_last_layer_only=True, |
| 106 | static_rnn_unroll_size=T, |
| 107 | ) |
| 108 | |
| 109 | if "dag" in args.implementation: |
| 110 | print("Using DAG net type") |
| 111 | model.net.Proto().type = 'dag' |
| 112 | model.net.Proto().num_workers = 4 |
| 113 | |
| 114 | elif args.implementation == "cudnn": |
| 115 | # We need to feed a placeholder input so that RecurrentInitOp |
| 116 | # can infer the dimensions. |
| 117 | init_blobs = model.net.AddExternalInputs("hidden_init", "cell_init") |
| 118 | model.param_init_net.ConstantFill([], input_blob, shape=input_shape) |
| 119 | output, last_hidden, _ = rnn_cell.cudnn_LSTM( |
| 120 | model=model, |
| 121 | input_blob=input_blob, |
| 122 | initial_states=init_blobs, |
| 123 | dim_in=args.input_dim, |
| 124 | dim_out=args.hidden_dim, |
searching dependent graphs…