| from tensorboard.backend.event_processing import event_accumulator | |
| import os | |
| from shutil import copy2 | |
| from re import search as RSearch | |
| import pandas as pd | |
| from ast import literal_eval as LEval | |
| weights_dir = 'logs/weights/' | |
| def find_biggest_tensorboard(tensordir): | |
| try: | |
| files = [f for f in os.listdir(tensordir) if f.endswith('.0')] | |
| if not files: | |
| print("No files with the '.0' extension found!") | |
| return | |
| max_size = 0 | |
| biggest_file = "" | |
| for file in files: | |
| file_path = os.path.join(tensordir, file) | |
| if os.path.isfile(file_path): | |
| file_size = os.path.getsize(file_path) | |
| if file_size > max_size: | |
| max_size = file_size | |
| biggest_file = file | |
| return biggest_file | |
| except FileNotFoundError: | |
| print("Couldn't find your model!") | |
| return | |
| def main(model_name, save_freq, lastmdls): | |
| global lowestval_weight_dir, scl | |
| tensordir = os.path.join('logs', model_name) | |
| lowestval_weight_dir = os.path.join(tensordir, "lowestvals") | |
| latest_file = find_biggest_tensorboard(tensordir) | |
| if latest_file is None: | |
| print("Couldn't find a valid tensorboard file!") | |
| return | |
| tfile = os.path.join(tensordir, latest_file) | |
| ea = event_accumulator.EventAccumulator(tfile, | |
| size_guidance={ | |
| event_accumulator.COMPRESSED_HISTOGRAMS: 500, | |
| event_accumulator.IMAGES: 4, | |
| event_accumulator.AUDIO: 4, | |
| event_accumulator.SCALARS: 0, | |
| event_accumulator.HISTOGRAMS: 1, | |
| }) | |
| ea.Reload() | |
| ea.Tags() | |
| scl = ea.Scalars('loss/g/total') | |
| listwstep = {} | |
| for val in scl: | |
| if (val.step // save_freq) * save_freq in [val.step for val in scl]: | |
| listwstep[float(val.value)] = (val.step // save_freq) * save_freq | |
| lowest_vals = sorted(listwstep.keys())[:lastmdls] | |
| sorted_dict = {value: step for value, step in listwstep.items() if value in lowest_vals} | |
| return sorted_dict | |
| def selectweights(model_name, file_dict, weights_dir, lowestval_weight_dir): | |
| os.makedirs(lowestval_weight_dir, exist_ok=True) | |
| logdir = [] | |
| files = [] | |
| lbldict = { | |
| 'Values': {}, | |
| 'Names': {} | |
| } | |
| weights_dir_path = os.path.join(weights_dir, "") | |
| low_val_path = os.path.join(os.getcwd(), os.path.join(lowestval_weight_dir, "")) | |
| try: | |
| file_dict = LEval(file_dict) | |
| except Exception as e: | |
| print(f"Error! {e}") | |
| return f"Couldn't load tensorboard file! {e}" | |
| weights = [f for f in os.scandir(weights_dir)] | |
| for key, value in file_dict.items(): | |
| pattern = fr"^{model_name}_.*_s{value}\.pth$" | |
| matching_weights = [f.name for f in weights if f.is_file() and RSearch(pattern, f.name)] | |
| for weight in matching_weights: | |
| source_path = weights_dir_path + weight | |
| destination_path = os.path.join(lowestval_weight_dir, weight) | |
| copy2(source_path, destination_path) | |
| logdir.append(f"File = {weight} Value: {key}, Step: {value}") | |
| lbldict['Names'][weight] = weight | |
| lbldict['Values'][weight] = key | |
| files.append(low_val_path + weight) | |
| print(f"File = {weight} Value: {key}, Step: {value}") | |
| yield ('\n'.join(logdir), files, pd.DataFrame(lbldict)) | |
| return ''.join(logdir), files, pd.DataFrame(lbldict) | |
| if __name__ == "__main__": | |
| model = str(input("Enter the name of the model: ")) | |
| sav_freq = int(input("Enter save frequency of the model: ")) | |
| ds = main(model, sav_freq) | |
| if ds: selectweights(model, ds, weights_dir, lowestval_weight_dir) | |