Paolify/RVC_v2

1from tensorboard.backend.event_processing import event_accumulator2 3import os4from shutil import copy25from re import search as RSearch6import pandas as pd7from ast import literal_eval as LEval8 9weights_dir = 'weights/'10 11def find_biggest_tensorboard(tensordir):12    try:13        files = [f for f in os.listdir(tensordir) if f.endswith('.0')]14        if not files:15            print("No files with the '.0' extension found!")16            return17 18        max_size = 019        biggest_file = ""20 21        for file in files:22            file_path = os.path.join(tensordir, file)23            if os.path.isfile(file_path):24                file_size = os.path.getsize(file_path)25                if file_size > max_size:26                    max_size = file_size27                    biggest_file = file28 29        return biggest_file30 31    except FileNotFoundError:32        print("Couldn't find your model!")33        return34 35def main(model_name, save_freq, lastmdls):36    global lowestval_weight_dir, scl37 38    tensordir = os.path.join('logs', model_name)39    lowestval_weight_dir = os.path.join(tensordir, "lowestvals")40    41    latest_file = find_biggest_tensorboard(tensordir)42    43    if latest_file is None:44        print("Couldn't find a valid tensorboard file!")45        return46    47    tfile = os.path.join(tensordir, latest_file)48    49    ea = event_accumulator.EventAccumulator(tfile,50        size_guidance={51        event_accumulator.COMPRESSED_HISTOGRAMS: 500,52        event_accumulator.IMAGES: 4,53        event_accumulator.AUDIO: 4,54        event_accumulator.SCALARS: 0,55        event_accumulator.HISTOGRAMS: 1,56    })57 58    ea.Reload()59    ea.Tags()60 61    scl = ea.Scalars('loss/g/total')62 63    listwstep = {}64    65    for val in scl:66        if (val.step // save_freq) * save_freq in [val.step for val in scl]:67            listwstep[float(val.value)] = (val.step // save_freq) * save_freq68 69    lowest_vals = sorted(listwstep.keys())[:lastmdls]70 71    sorted_dict = {value: step for value, step in listwstep.items() if value in lowest_vals}72    73    return sorted_dict74 75def selectweights(model_name, file_dict, weights_dir, lowestval_weight_dir):76    os.makedirs(lowestval_weight_dir, exist_ok=True)77    logdir = []78    files = []79    lbldict = {80        'Values': {},81        'Names': {}82    }83    weights_dir_path = os.path.join(weights_dir, "")84    low_val_path = os.path.join(os.getcwd(), os.path.join(lowestval_weight_dir, ""))85    86    try:87        file_dict = LEval(file_dict)88    except Exception as e: 89        print(f"Error! {e}")90        return f"Couldn't load tensorboard file! {e}"91    92    weights = [f for f in os.scandir(weights_dir)]93    for key, value in file_dict.items():94        pattern = fr"^{model_name}_.*_s{value}\.pth$"95        matching_weights = [f.name for f in weights if f.is_file() and RSearch(pattern, f.name)]96        for weight in matching_weights:97            source_path = weights_dir_path + weight98            destination_path = os.path.join(lowestval_weight_dir, weight)99            100            copy2(source_path, destination_path)101 102            logdir.append(f"File = {weight} Value: {key}, Step: {value}")103 104            lbldict['Names'][weight] = weight105            lbldict['Values'][weight] = key106 107            files.append(low_val_path + weight)108 109            print(f"File = {weight} Value: {key}, Step: {value}")110 111            yield ('\n'.join(logdir), files, pd.DataFrame(lbldict))112            113 114    return ''.join(logdir), files, pd.DataFrame(lbldict)115    116 117if __name__ == "__main__":118    model = str(input("Enter the name of the model: "))119    sav_freq = int(input("Enter save frequency of the model: "))120    ds = main(model, sav_freq)121    122    if ds: selectweights(model, ds, weights_dir, lowestval_weight_dir)123