| import logging |
| import os |
| import numpy as np |
| import copy |
| from utils.metrics import clustering_score |
| from sklearn.metrics import confusion_matrix |
| from keras.models import Model |
| from keras.optimizers import SGD |
| from tqdm import trange |
| from configs.base import ParamManager |
| from utils.functions import set_seed |
| from backbones.sae import get_sae, ClusteringLayer |
| from sklearn.cluster import KMeans |
|
|
| def target_distribution(q): |
| weight = q ** 2 / q.sum(0) |
| return (weight.T / weight.sum(1)).T |
|
|
| class DECManager: |
| |
| def __init__(self, args, data, model, logger_name = 'Discovery'): |
| |
| self.logger = logging.getLogger(logger_name) |
| |
|
|
| self.tfidf_train, self.tfidf_test = data.dataloader.tfidf_train, data.dataloader.tfidf_test |
| self.num_labels = data.num_labels |
| self.test_y = data.dataloader.test_true_labels |
|
|
| self.init_sae(args, data, model) |
| set_seed(args.seed) |
| if args.train: |
| self.model, self.y_pred_init = self.init_model(args) |
|
|
| else: |
| clustering_layer = ClusteringLayer(self.num_labels, name='clustering')(self.sae.layers[3].output) |
| self.model = Model(inputs=self.sae.input, outputs = [clustering_layer, self.sae.output]) |
| |
| save_path = os.path.join(args.model_output_dir, args.model_name) |
| self.logger.info('Loading models from %s' % save_path) |
| self.model.load_weights(save_path) |
|
|
| def init_sae(self, args, data, model): |
| self.sae = model.set_model(args, data, 'sae') |
| self.sae_feats_path = os.path.join(args.model_output_dir, 'SAE.h5') |
| |
| if os.path.exists(self.sae_feats_path): |
| self.logger.info('Loading SAE features from %s' % self.sae_feats_path) |
| self.sae.load_weights(self.sae_feats_path) |
| else: |
| self.logger.info('SAE (emb) training start...') |
| self.sae.fit(self.tfidf_train, self.tfidf_train, epochs = args.num_train_epochs_SAE, batch_size = args.SAE_batch_size, shuffle=True, |
| validation_data=(self.tfidf_test, self.tfidf_test), verbose=1) |
| self.logger.info('SAE (emb) training finished...') |
|
|
| if args.save_model: |
|
|
| save_path = os.path.join(args.model_output_dir, 'SAE.h5') |
| self.logger.info('Save models at %s', str(save_path)) |
| self.sae.save_weights(save_path) |
|
|
| def init_model(self, args): |
|
|
| sae_emb_train, sae_emb_test = get_sae(args, self.sae, self.tfidf_train, self.tfidf_test) |
| clustering_layer = ClusteringLayer(self.num_labels, name='clustering')(self.sae.layers[3].output) |
| model = Model(inputs=self.sae.input, outputs = clustering_layer) |
| model.compile(optimizer=SGD(args.lr, args.momentum), loss='kld') |
|
|
| km = KMeans(n_clusters=self.num_labels, n_init=20, n_jobs=-1, random_state=args.seed) |
| y_pred = km.fit_predict(sae_emb_train) |
| y_pred_last = np.copy(y_pred) |
| model.get_layer(name='clustering').set_weights([km.cluster_centers_]) |
|
|
| return model, y_pred_last |
|
|
| def train(self, args, data): |
|
|
| self.logger.info('DEC training starts...') |
| index = 0 |
| loss = 0 |
| index_array = np.arange(self.tfidf_train.shape[0]) |
| y_pred_last = self.y_pred_init |
|
|
| for epoch in trange(int(args.num_train_epochs_DEC), desc="Epoch"): |
|
|
| if epoch % args.update_interval == 0: |
|
|
| q = self.model.predict(self.tfidf_train, verbose=0) |
| p = target_distribution(q) |
|
|
| y_pred = q.argmax(1) |
| delta_label = np.sum(y_pred != y_pred_last).astype(np.float32) / y_pred.shape[0] |
| y_pred_last = np.copy(y_pred) |
| |
| if epoch > 0: |
| |
| self.logger.info("***** Epoch: %s*****", str(epoch + 1)) |
| self.logger.info('Training Loss: %f', np.round(loss, 5)) |
| self.logger.info('Delta Label: %f', delta_label) |
|
|
| if delta_label < args.tol: |
| self.logger.info('delta_label %s < %f', delta_label, args.tol) |
| self.logger.info('Reached tolerance threshold. Stop training.') |
| break |
|
|
| idx = index_array[index * args.DEC_batch_size: min((index + 1) * args.DEC_batch_size, self.tfidf_train.shape[0])] |
| loss = self.model.train_on_batch(x = self.tfidf_train[idx], y = p[idx]) |
| index = index + 1 if (index + 1) * args.DEC_batch_size <= self.tfidf_train.shape[0] else 0 |
|
|
| self.logger.info('DEC training finished...') |
|
|
| if args.save_model: |
| save_path = os.path.join(args.model_output_dir, args.model_name) |
| self.model.save_weights(save_path) |
|
|
|
|
| def test(self, args, data, show=False): |
|
|
| q = self.model.predict(self.tfidf_test, verbose = 0) |
| y_pred = q.argmax(1) |
| y_true = self.test_y |
|
|
| test_results = clustering_score(y_true, y_pred) |
| cm = confusion_matrix(y_true,y_pred) |
| |
| if show: |
| self.logger.info |
| self.logger.info("***** Test: Confusion Matrix *****") |
| self.logger.info("%s", str(cm)) |
| self.logger.info("***** Test results *****") |
| |
| for key in sorted(test_results.keys()): |
| self.logger.info(" %s = %s", key, str(test_results[key])) |
|
|
| test_results['y_true'] = y_true |
| test_results['y_pred'] = y_pred |
|
|
| return test_results |
| |
|
|
