Function train

scripts/release_notes/classifier.py:326–382

(save_path: Path, data_folder: Path, regen_data: bool, resample: bool)

def train(save_path: Path, data_folder: Path, regen_data: bool, resample: bool):
    train_data, val_data = get_train_val_data(data_folder, regen_data)
    train_zip_list = get_title_files_author_categories_zip_list(train_data)
    val_zip_list = get_title_files_author_categories_zip_list(val_data)

    classifier_config = CategoryConfig(common.categories)
    author_map = get_author_map(data_folder, regen_data)
    file_map = get_file_map(data_folder, regen_data)
    commit_classifier = CommitClassifier(
        XLMR_BASE, author_map, file_map, classifier_config
    ).to(device)

    # Lets train this bag of bits
    class_weights = gen_class_weights(train_zip_list)
    loss = torch.nn.CrossEntropyLoss(weight=class_weights)
    optimizer = torch.optim.Adam(commit_classifier.parameters(), lr=3e-3)

    num_epochs = 25
    batch_size = 256

    if resample:
        # Lets not use this
        train_zip_list = balance_dataset(train_zip_list)
    data_size = len(train_zip_list)

    print(f"Training on {data_size} examples.")
    # We can fit all of val into one batch
    val_batch = generate_batch(val_zip_list)

    for i in tqdm(range(num_epochs), desc="Epochs"):
        start = 0
        random.shuffle(train_zip_list)
        while start < data_size:
            end = start + batch_size
            # make the last batch bigger if needed
            if end > data_size:
                end = data_size
            train_batch = train_zip_list[start:end]
            train_batch = generate_batch(train_batch)
            l = train_step(train_batch, commit_classifier, optimizer, loss)
            start = end

        val_l = eval_step(val_batch, commit_classifier, loss)
        tqdm.write(
            f"Finished epoch {i} with a train loss of: {l.item()} and a val_loss of: {val_l.item()}"
        )

    with torch.no_grad():
        commit_classifier.eval()
        val_inpts, val_targets = val_batch
        val_output = commit_classifier(val_inpts)
        val_preds = torch.argmax(val_output, dim=1)
        val_acc = torch.sum(val_preds == val_targets).item() / len(val_preds)
        print(f"Final Validation accuracy is {val_acc}")

    print(f"Jobs done! Saving to {save_path}")
    torch.save(commit_classifier.state_dict(), save_path)