ensembling.py (11106B)
1 from typing import Any, Dict, List, Tuple 2 import argparse 3 import collections 4 import json 5 import os 6 import re 7 import sys 8 import pathlib 9 import pickle 10 11 import tqdm 12 import numpy 13 import sklearn.metrics 14 import sklearn.linear_model 15 import xgboost 16 17 import mmsrl.utils 18 19 20 def build_model_paths(paths): 21 model_paths = set() 22 for path in paths: 23 for filename in os.listdir(path): 24 modelname = filename.replace("val", "{}").replace("test", "{}") 25 if modelname.count("{}") == 1: 26 model_paths.add(path / modelname) 27 elif modelname.count("{}") > 1: 28 print(f"\033[31;1mWARNING: ignoring file with several \"val\" and \"test\" in filename: {filename}.\033[0m") 29 return model_paths 30 31 32 def read_prediction_file(path): 33 try: 34 data = collections.defaultdict(dict) 35 with open(path, "rb") as file: 36 try: 37 while True: 38 image, entities = pickle.load(file) 39 for entity, prediction in entities.items(): 40 data[image][entity] = prediction.numpy() 41 except EOFError: 42 return data 43 except FileNotFoundError: 44 return {} 45 46 47 def read_predictions(paths): 48 data = {} 49 for path in tqdm.tqdm(paths, desc="Loading predictions", leave=False): 50 prediction = str(path) 51 data[prediction] = {} 52 for split in ["val", "test"]: 53 path = prediction.replace("{}", split) 54 data[prediction][split] = read_prediction_file(path) 55 return data 56 57 58 def read_data(split): 59 data_path = pathlib.Path(os.environ["DATA_PATH"]) 60 path = data_path / "mmsrl" / "all" / "annotations" / f"{split}.jsonl" 61 with path.open("r") as file: 62 return list(map(json.loads, file)) 63 64 65 def write_prediction_file(path, data): 66 with open(path, "wb") as file: 67 for image, predictions in data.items(): 68 pickle.dump((image, predictions), file) 69 70 71 def get_sample_entities(sample): 72 if "entity_list" in sample: 73 return sample["entity_list"] 74 else: 75 return sum((sample[label] for label in mmsrl.utils.LABELS), []) 76 77 78 def models_to_tensors(data, models): 79 tensors = {} 80 for sample in data: 81 image = sample["image"] 82 tensors[image] = {} 83 for entity in set(get_sample_entities(sample)): 84 if not entity: 85 continue 86 values = [] 87 for model in models: 88 values.append(model[image][entity]) 89 tensors[image][entity] = numpy.stack(values) 90 return tensors 91 92 93 def model_is_valid(model, dataset): 94 try: 95 for sample in dataset: 96 image = sample["image"] 97 for entity in get_sample_entities(sample): 98 if entity: 99 pred = model[image][entity] 100 if not numpy.isfinite(pred).all(): 101 return False 102 except KeyError: 103 return False 104 return True 105 106 107 def compute_score(predictions, truth, final, split): 108 y_pred = [] 109 y_truth = [] 110 for sample in truth: 111 image = sample["image"] 112 for ilabel, label in enumerate(mmsrl.utils.LABELS): 113 for entity in sample[label]: 114 if entity: 115 y_pred.append(predictions[image][entity].argmax().item()) 116 else: 117 y_pred.append(3) 118 y_truth.append(ilabel) 119 120 f1 = sklearn.metrics.f1_score(y_truth, y_pred, average="macro", zero_division=0) 121 class_precision, class_recall, class_f1, _ = sklearn.metrics.precision_recall_fscore_support(y_truth, y_pred, zero_division=0) 122 class_scores = ", ".join(f"{label}[F{class_f1[i]:4.2f}/P{class_precision[i]:4.2f}/R{class_recall[i]:4.2f}]" for i, label in enumerate(mmsrl.utils.LABELS)) 123 color = "32" if final else "33" 124 return f1, f"{split} Macro-F1: \033[{color};1m{f1:5.3f}\033[0m, {class_scores}" 125 126 127 def models_to_xy(models, valid): 128 nmodels = len(models) 129 x = [] 130 y = [] 131 order = [] 132 for sample in valid: 133 image = sample["image"] 134 for ilabel, label in enumerate(mmsrl.utils.LABELS): 135 for entity in sample[label]: 136 if not entity: 137 continue 138 row = numpy.empty((nmodels,4), dtype=numpy.float32) 139 for im, model in enumerate(models): 140 row[im] = model[image][entity] 141 x.append(row) 142 y.append(ilabel) 143 order.append((image, entity)) 144 return numpy.stack(x), numpy.array(y, dtype=numpy.int64), order 145 146 147 def build_ensemble(config, valid, models): 148 tmodels = models_to_tensors(valid, models) 149 if config.aggregator in ["linear", "xgboost"]: 150 x, y, order = models_to_xy(models, valid) 151 152 if config.aggregator in "linear": 153 if config.full_matrix: 154 linear = sklearn.linear_model.LogisticRegression(penalty="none", class_weight="balanced", max_iter=1000) 155 linear.fit(x.reshape(x.shape[0], -1), y) 156 else: 157 nsample = y.shape[0] 158 label_weight = nsample / 4 / numpy.bincount(y) 159 nmodels = len(models) 160 x = x.transpose(0, 2, 1) 161 x = x.reshape(-1, nmodels) 162 y = y.repeat(4) 163 sample_weight = label_weight[y] 164 y = (y.reshape(nsample, 4) == numpy.arange(4)).flatten() 165 166 linear = sklearn.linear_model.LogisticRegression(penalty="none", max_iter=1000) 167 linear.fit(x, y, sample_weight=sample_weight) 168 ensemble = linear.coef_ 169 elif config.aggregator == "xgboost": 170 nsample = y.shape[0] 171 label_weight = nsample / 4 / numpy.bincount(y) 172 param = {'objective': 'multi:softmax', 'num_class': 4} 173 param.update(eval(config.xgboost_params)) 174 print(f"\033[33mXGBoost parameters: {param}\033[0m") 175 ensemble = xgboost.XGBRegressor(**param) 176 ensemble.fit(x.reshape(2067, -1), y, sample_weight=label_weight[y]) 177 if config.save_xgboost: 178 ensemble.save_model(str(config.output / get_output_name(config)) + ".xgboost.json") 179 with open(str(config.output / get_output_name(config)) + ".xgboost.map", "wb") as outfile: 180 pickle.dump(order, outfile) 181 else: 182 ensemble = None 183 184 prediction = ensemble_predict(config, valid, ensemble, models) 185 print(compute_score(prediction, valid, True, "VALID")[1]) 186 return ensemble 187 188 189 def ensemble_predict_sample(config, ensemble, models): 190 if config.aggregator == "mean": 191 return models.mean(0) 192 elif config.aggregator == "max": 193 return models.max(0)[0] 194 elif config.aggregator == "min": 195 return models.min(0)[0] 196 elif config.aggregator == "median": 197 models.sort(0) 198 return models[models.shape[0]//2, :] 199 elif config.aggregator == "linear": 200 return ensemble.dot(models.flatten() if config.full_matrix else models) 201 elif config.aggregator == "xgboost": 202 probabilities = numpy.zeros(4, dtype=numpy.float32) 203 prediction = ensemble.predict(models.reshape(1, -1)) 204 probabilities[int(prediction)] = 1 205 return probabilities 206 else: 207 raise RuntimeError(f"Unknown aggregator {config.aggregator}") 208 209 210 def ensemble_predict(config, test, ensemble, models): 211 tmodels = models_to_tensors(test, models) 212 output = {} 213 for sample in test: 214 image = sample["image"] 215 output[image] = {} 216 for entity in set(get_sample_entities(sample)): 217 if entity: 218 output[image][entity] = ensemble_predict_sample(config, ensemble, tmodels[image][entity]) 219 else: 220 output[image][entity] = numpy.array([0, 0, 0, 1], dtype=numpy.float32) 221 return output 222 223 224 def filter_models(config: mmsrl.utils.dotdict, models, val_data, test_data): 225 useme = {} 226 for model_name, model in models.items(): 227 if not (model_is_valid(model["val"], val_data) and model_is_valid(model["test"], test_data)): 228 continue 229 f1, msg = compute_score(model["val"], val_data, False, "VALID") 230 if f1 > config.cut_off: 231 useme[model_name] = model 232 if config.show_all: 233 print(f"\033[33m{model_name}\033[0m") 234 print(msg) 235 return useme 236 237 238 def get_output_name(config: mmsrl.utils.dotdict) -> str: 239 name = f"ensemble_{config.aggregator}_" 240 name += str(config.cut_off).split('.')[-1] 241 if config.aggregator == "linear" and config.full_matrix: 242 name += "_fullmatrix" 243 if config.aggregator == "xgboost": 244 name += re.sub("[^a-zA-Z0-9.]+", "_", config.xgboost_params).rstrip('_') 245 return f"{name}.pkl" 246 247 248 def perform_ensembling(config: mmsrl.utils.dotdict, models, val_data, test_data): 249 models_filtered = filter_models(config, models, val_data, test_data) 250 models_order = list(models_filtered.keys()) 251 models = [models[key] for key in models_order] 252 ensemble = build_ensemble(config, val_data, [model["val"] for model in models]) 253 prediction = ensemble_predict(config, test_data, ensemble, [model["test"] for model in models]) 254 outfile = config.output / get_output_name(config) 255 print(f"Writing test prediction to \033[1m{outfile}\033[0m") 256 write_prediction_file(outfile, prediction) 257 print(compute_score(prediction, test_data, True, "TEST ")[1]) 258 259 260 def main(config: mmsrl.utils.dotdict) -> None: 261 val_data = read_data("val") 262 test_data = read_data("test") 263 models_path = build_model_paths(config.prediction) 264 models = read_predictions(models_path) 265 if config.cut_off==0: 266 for cut in range(31, 45, 2): 267 print(cut/100) 268 config.cut_off = cut /100 269 perform_ensembling(config, models, val_data, test_data) 270 else: 271 perform_ensembling(config, models, val_data, test_data) 272 273 274 if __name__ == "__main__": 275 parser = argparse.ArgumentParser() 276 parser.add_argument("output", type=pathlib.Path, help="Path were we write the test prediction output.") 277 parser.add_argument("prediction", type=pathlib.Path, nargs='+', help="Paths to directories containing prediction files, the filenames should contain \"val\" and \"test\" such that substituting one for other gives the corresponding predictions on the other split.") 278 parser.add_argument("--aggregator", type=str, default="mean", help="How to aggregate results.") 279 parser.add_argument("--full-matrix", action="store_true", help="When using logistic regression, learn a full 4*nmodel→4 matrix instead of a 4→1 one.") 280 parser.add_argument("--cut-off", type=float, default=0.3, help="Ignore models with a lower F1 on the validation set.") 281 parser.add_argument("--show-all", action="store_true", help="Show validation score of all (not cut-off) models.") 282 parser.add_argument("--save-xgboost", action="store_true", help="Save the xgboost model and mapping.") 283 parser.add_argument("--xgboost-params", type=str, default="{}", help="Should be eval()uable to a dict which is then added to xgboost parameters.") 284 args = mmsrl.utils.dotdict(vars(parser.parse_args())) 285 main(args)