def extract_embeddings(): parser = argparse.ArgumentParser(description='') parser.add_argument('database', type=str) parser.add_argument('collection', type=str) parser.add_argument('outfile', type=str) parser.add_argument('dimensions', type=int, default=100) parser.add_argument('epochs', type=int, default=10) args = parser.parse_args() model = EmbeddingModel(args.database, args.collection) model.train(args.dimensions, args.epochs) model.save(f'data/{args.outfile}')
args.device = torch.device('cuda') torch.cuda.manual_seed(np.random.randint(1, 10000)) torch.backends.cudnn.enabled = True args.classes = ["desert", "rainforest", "grassland", "tundra", "ocean"] if args.embedding_type == 'linear': model = EmbeddingModel(len(args.classes)) elif args.embedding_type == 'conv': model = ConvolutionalEmbeddingModel(len(args.classes)) else: print("Model type [{0}] not supported".format(args.embedding_type)) exit(1) if torch.cuda.is_available(): model = model.cuda() criterion = torch.nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=args.learning_rate) model = model.train() starting_epoch = 0 running_loss = 0.0 if len(args.model_checkpoint) > 0: checkpoint = torch.load(args.model_checkpoint) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) starting_epoch = checkpoint['epoch'] running_loss = checkpoint['running_loss'] train_set, train_labels, test_set, test_labels, classes = generateData(args.corpus_file, args.classes, args.train_split_percentage, args.load_embedding_from_file, args.save_embedding_dict) print(len(train_set)) print('Training model...') print('Starting from epoch %d' % (starting_epoch + 1)) for epoch in progressbar(range(starting_epoch, args.epochs)): for i in range(len(train_set)): inputs = train_set[i]
help='size of skip window') parser.add_argument('--batch_size', type=int, dest='batch_size', help='size of batch') parser.add_argument('--num_steps', type=int, dest='num_steps', help='total iterations') parser.add_argument('--display_steps', type=int, dest='display_steps', help='display steps') parser.add_argument('--e_steps', type=int, dest='e_steps', help='e steps') parser.add_argument('--learning_rate', type=float, dest='learning_rate', help='learning rate') c = Config() parser.parse_args(namespace=c) return c if __name__ == '__main__': conf = parse_args() df = DataFactory(conf) df.load_data() t0 = time.time() with tf.Session() as sess: m = EmbeddingModel(conf) init_op = tf.initialize_all_variables() sess.run(init_op) m.train(sess, df) print 'Done train model, cost time: %0.3fs' % (time.time() - t0)