Hilly12 / prifair

PriFair

A light-weight library for training and evaluating fair and privacy-preserving machine learning models. Built on top of PyTorch and Opacus.

Installation

pip install git+https://github.com/Hilly12/prifair.git

Usage

Training a DP-SGD model with reweighing.

import prifair as pf

# Setup

weights = pf.core.reweigh(train_labels)

train_loader, model, optimizer, accountant = pf.core.setup_weighted_dpsgd(
    data_loader=train_loader,
    model=model,
    optimizer=optimizer,
    weights=weights,
    target_epsilon=target_epsilon,
    target_delta=target_delta,
    max_grad_norm=max_grad_norm,
    epochs=epochs,
)

# Standard Opacus Training

for epoch in range(epochs):

    with BatchMemoryManager(
        data_loader=train_loader,
        max_physical_batch_size=max_physical_batch_size,
        optimizer=optimizer,
    ) as memory_safe_data_loader:

        for images, target in tqdm(memory_safe_data_loader):
            optimizer.zero_grad()
            images = images.to(device)
            target = target.to(device)

            output = model(images)
            loss = criterion(output, target)

            loss.backward()
            optimizer.step()

epsilon = accountant.get_epsilon(target_delta)