# pylint:skip-file from __future__ import print_function import logging import sys, random, time, math sys.path.insert(0, "../../python") import mxnet as mx import numpy as np from collections import namedtuple from nce import * from operator import itemgetter from optparse import OptionParser LSTMState = namedtuple("LSTMState", ["c", "h"]) LSTMParam = namedtuple("LSTMParam", ["i2h_weight", "i2h_bias", "h2h_weight", "h2h_bias"]) LSTMModel = namedtuple("LSTMModel", ["rnn_exec", "symbol", "init_states", "last_states", "seq_data", "seq_labels", "seq_outputs", "param_blocks"]) def lstm(num_hidden, indata, prev_state, param, seqidx, layeridx, dropout=0.): """LSTM Cell symbol""" if dropout > 0.: indata = mx.sym.Dropout(data=indata, p=dropout) i2h = mx.sym.FullyConnected(data=indata, weight=param.i2h_weight, bias=param.i2h_bias, num_hidden=num_hidden * 4, name="t%d_l%d_i2h" % (seqidx, layeridx)) h2h = mx.sym.FullyConnected(data=prev_state.h, weight=param.h2h_weight, bias=param.h2h_bias, num_hidden=num_hidden * 4, name="t%d_l%d_h2h" % (seqidx, layeridx)) gates = i2h + h2h slice_gates = mx.sym.SliceChannel(gates, num_outputs=4, name="t%d_l%d_slice" % (seqidx, layeridx)) in_gate = mx.sym.Activation(slice_gates[0], act_type="sigmoid") in_transform = mx.sym.Activation(slice_gates[1], act_type="tanh") forget_gate = mx.sym.Activation(slice_gates[2], act_type="sigmoid") out_gate = mx.sym.Activation(slice_gates[3], act_type="sigmoid") next_c = (forget_gate * prev_state.c) + (in_gate * in_transform) next_h = out_gate * mx.sym.Activation(next_c, act_type="tanh") return LSTMState(c=next_c, h=next_h) def get_net(vocab_size, seq_len, num_label, num_lstm_layer, num_hidden): param_cells = [] last_states = [] for i in range(num_lstm_layer): param_cells.append(LSTMParam(i2h_weight=mx.sym.Variable("l%d_i2h_weight" % i), i2h_bias=mx.sym.Variable("l%d_i2h_bias" % i), h2h_weight=mx.sym.Variable("l%d_h2h_weight" % i), h2h_bias=mx.sym.Variable("l%d_h2h_bias" % i))) state = LSTMState(c=mx.sym.Variable("l%d_init_c" % i), h=mx.sym.Variable("l%d_init_h" % i)) last_states.append(state) data = mx.sym.Variable('data') label = mx.sym.Variable('label') label_weight = mx.sym.Variable('label_weight') embed_weight = mx.sym.Variable('embed_weight') label_embed_weight = mx.sym.Variable('label_embed_weight') data_embed = mx.sym.Embedding(data = data, input_dim = vocab_size, weight = embed_weight, output_dim = 100, name = 'data_embed') datavec = mx.sym.SliceChannel(data = data_embed, num_outputs = seq_len, squeeze_axis = True, name = 'data_slice') labelvec = mx.sym.SliceChannel(data = label, num_outputs = seq_len, squeeze_axis = True, name = 'label_slice') labelweightvec = mx.sym.SliceChannel(data = label_weight, num_outputs = seq_len, squeeze_axis = True, name = 'label_weight_slice') probs = [] for seqidx in range(seq_len): hidden = datavec[seqidx] for i in range(num_lstm_layer): next_state = lstm(num_hidden, indata = hidden, prev_state = last_states[i], param = param_cells[i], seqidx = seqidx, layeridx = i) hidden = next_state.h last_states[i] = next_state probs.append(nce_loss(data = hidden, label = labelvec[seqidx], label_weight = labelweightvec[seqidx], embed_weight = label_embed_weight, vocab_size = vocab_size, num_hidden = 100, num_label = num_label)) return mx.sym.Group(probs) def load_data(name): buf = open(name).read() tks = buf.split(' ') vocab = {} freq = [0] data = [] for tk in tks: if len(tk) == 0: continue if tk not in vocab: vocab[tk] = len(vocab) + 1 freq.append(0) wid = vocab[tk] data.append(wid) freq[wid] += 1 negative = [] for i, v in enumerate(freq): if i == 0 or v < 5: continue v = int(math.pow(v * 1.0, 0.75)) negative += [i for _ in range(v)] return data, negative, vocab, freq class SimpleBatch(object): def __init__(self, data_names, data, label_names, label): self.data = data self.label = label self.data_names = data_names self.label_names = label_names @property def provide_data(self): return [(n, x.shape) for n, x in zip(self.data_names, self.data)] @property def provide_label(self): return [(n, x.shape) for n, x in zip(self.label_names, self.label)] class DataIter(mx.io.DataIter): def __init__(self, name, batch_size, seq_len, num_label, init_states): super(DataIter, self).__init__() self.batch_size = batch_size self.data, self.negative, self.vocab, self.freq = load_data(name) self.vocab_size = 1 + len(self.vocab) print(self.vocab_size) self.seq_len = seq_len self.num_label = num_label self.init_states = init_states self.init_state_names = [x[0] for x in self.init_states] self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states] self.provide_data = [('data', (batch_size, seq_len))] + init_states self.provide_label = [('label', (self.batch_size, seq_len, num_label)), ('label_weight', (self.batch_size, seq_len, num_label))] def sample_ne(self): return self.negative[random.randint(0, len(self.negative) - 1)] def __iter__(self): print('begin') batch_data = [] batch_label = [] batch_label_weight = [] for i in range(0, len(self.data) - self.seq_len - 1, self.seq_len): data = self.data[i: i+self.seq_len] label = [[self.data[i+k+1]] \ + [self.sample_ne() for _ in range(self.num_label-1)]\ for k in range(self.seq_len)] label_weight = [[1.0] \ + [0.0 for _ in range(self.num_label-1)]\ for k in range(self.seq_len)] batch_data.append(data) batch_label.append(label) batch_label_weight.append(label_weight) if len(batch_data) == self.batch_size: data_all = [mx.nd.array(batch_data)] + self.init_state_arrays label_all = [mx.nd.array(batch_label), mx.nd.array(batch_label_weight)] data_names = ['data'] + self.init_state_names label_names = ['label', 'label_weight'] batch_data = [] batch_label = [] batch_label_weight = [] yield SimpleBatch(data_names, data_all, label_names, label_all) def reset(self): pass if __name__ == '__main__': head = '%(asctime)-15s %(message)s' logging.basicConfig(level=logging.DEBUG, format=head) parser = OptionParser() parser.add_option("-g", "--gpu", action = "store_true", dest = "gpu", default = False, help = "use gpu") batch_size = 1024 seq_len = 5 num_label = 6 num_lstm_layer = 2 num_hidden = 100 init_c = [('l%d_init_c'%l, (batch_size, num_hidden)) for l in range(num_lstm_layer)] init_h = [('l%d_init_h'%l, (batch_size, num_hidden)) for l in range(num_lstm_layer)] init_states = init_c + init_h data_train = DataIter("./data/text8", batch_size, seq_len, num_label, init_states) network = get_net(data_train.vocab_size, seq_len, num_label, num_lstm_layer, num_hidden) options, args = parser.parse_args() devs = mx.cpu() if options.gpu == True: devs = mx.gpu() model = mx.model.FeedForward(ctx = devs, symbol = network, num_epoch = 20, learning_rate = 0.3, momentum = 0.9, wd = 0.0000, initializer=mx.init.Xavier(factor_type="in", magnitude=2.34)) metric = NceLSTMAuc() model.fit(X = data_train, eval_metric = metric, batch_end_callback = mx.callback.Speedometer(batch_size, 50),)