| 36 | |
| 37 | |
| 38 | class CharRNN: |
| 39 | def __init__(self, args): |
| 40 | self.seq_length = args.seq_length |
| 41 | self.batch_size = args.batch_size |
| 42 | self.iters_to_report = args.iters_to_report |
| 43 | self.hidden_size = args.hidden_size |
| 44 | |
| 45 | with open(args.train_data) as f: |
| 46 | self.text = f.read() |
| 47 | |
| 48 | self.vocab = list(set(self.text)) |
| 49 | self.char_to_idx = {ch: idx for idx, ch in enumerate(self.vocab)} |
| 50 | self.idx_to_char = {idx: ch for idx, ch in enumerate(self.vocab)} |
| 51 | self.D = len(self.char_to_idx) |
| 52 | |
| 53 | print("Input has {} characters. Total input size: {}".format( |
| 54 | len(self.vocab), len(self.text))) |
| 55 | |
| 56 | def CreateModel(self): |
| 57 | log.debug("Start training") |
| 58 | model = model_helper.ModelHelper(name="char_rnn") |
| 59 | |
| 60 | input_blob, seq_lengths, hidden_init, cell_init, target = \ |
| 61 | model.net.AddExternalInputs( |
| 62 | 'input_blob', |
| 63 | 'seq_lengths', |
| 64 | 'hidden_init', |
| 65 | 'cell_init', |
| 66 | 'target', |
| 67 | ) |
| 68 | |
| 69 | hidden_output_all, self.hidden_output, _, self.cell_state = LSTM( |
| 70 | model, input_blob, seq_lengths, (hidden_init, cell_init), |
| 71 | self.D, self.hidden_size, scope="LSTM") |
| 72 | output = brew.fc( |
| 73 | model, |
| 74 | hidden_output_all, |
| 75 | None, |
| 76 | dim_in=self.hidden_size, |
| 77 | dim_out=self.D, |
| 78 | axis=2 |
| 79 | ) |
| 80 | |
| 81 | # axis is 2 as first two are T (time) and N (batch size). |
| 82 | # We treat them as one big batch of size T * N |
| 83 | softmax = model.net.Softmax(output, 'softmax', axis=2) |
| 84 | |
| 85 | softmax_reshaped, _ = model.net.Reshape( |
| 86 | softmax, ['softmax_reshaped', '_'], shape=[-1, self.D]) |
| 87 | |
| 88 | # Create a copy of the current net. We will use it on the forward |
| 89 | # pass where we don't need loss and backward operators |
| 90 | self.forward_net = core.Net(model.net.Proto()) |
| 91 | |
| 92 | xent = model.net.LabelCrossEntropy([softmax_reshaped, target], 'xent') |
| 93 | # Loss is average both across batch and through time |
| 94 | # Thats why the learning rate below is multiplied by self.seq_length |
| 95 | loss = model.net.AveragedLoss(xent, 'loss') |
no outgoing calls
no test coverage detected
searching dependent graphs…