Package: glm-saga
SAGA-based GPU solver for elastic net problemsA package for fitting sparse linear models at deep learning scales. This work was initially created and described in our paper,"Leveraging Sparse Linear Layers for Debuggable Deep Networks"withEric Wong,Shibani SanturkarandAleksander Madry. The main repository for this paper can be foundhere.This package implements a SAGA-based solver in PyTorch for fitting sparse linear models with elastic net regularization. It combines the path algorithm used byglmnetwith aminibatch variant of the SAGA algorithm, which allows solving the elastic net at ImageNet scales, which coordinate descent-based elastic net solvers struggle with.CitationIf you find this solver to be useful in your work, consider citing our paper:@article{wong2021leveraging,title={Leveraging Sparse Linear Layers for Debuggable Deep Networks},author={Wong, Eric and Santurkar, Shibani and M{\k{a}}dry, Aleksander},journal={arXiv preprint arXiv:2105.04857},year={2021}}InstallationThis package is on PyPI. Install it withpip install glm_saga. The only requirement is PyTorch. Alternatively, the entire solver is implemented inglm_saga/elasticnet.pyand can be copied locally if desired.Usage and documentationThe following function is the main interface which can be used to fit a sequence of sparse linear models. A barebones example which fits a sparse linear model on top of a ResNet18 can be found inresnet18_example.py.def glm_saga(linear, loader, max_lr, nepochs, alpha, 
             table_device=None, precompute=None, group=False, 
             verbose=None, state=None, n_ex=None, n_classes=None, 
             tol=1e-4, epsilon=0.001, k=100, checkpoint=None, 
             solver='saga', do_zero=True, lr_decay_factor=1, metadata=None, 
             val_loader=None, test_loader=None, lookbehind=None, 
             family='multinomial'):Required argumentslinear: a PyTorchnn.Linearmodule which the solver initializes from (initialize this to zero)loader: a dataloader which returns examples in the form(X,y,i)whereXis a batch of features,yis a batch of labels, andiis a batch of indices which uniquely identify each example.Important: the features must be normalized (zero mean and unit variance) and the index is necessary for the solver. Optionally, the dataloader can also return(X,y,i,w)wherewis the sample weight.max_lr: the starting learning rate to use for the SAGA solver at the starting regularizationnepochs: the maximum number of epochs to run the SAGA solver for each step of regularizationalpha: a hyperparameter for elastic net regularization which controls the tradeoff between L1 and L2 regularization (typically taken to be 0.8 or 0.99).alpha=1corresponds to only L1 regularization, whereasalpha=0corresponds to only L2 regularization.Optional argumentstable_device=None: if specified, manually stores the SAGA gradient table on the specified device (otherwise, defaults to the device fo the given model). Useful for reducing memory usage.precompute=None: If specified, passes each example from the loader throughprecompute, which is assumed to be a PyTorchnn.Module. This can be used to normalize the data if the dataloaders are not already normalized.group=False: If true, use the grouped LASSO where groups are all parameters for a given feature. If false, use standard LASSO.verbose=None: If set to an integer, print the status of the inner GLM solver everyverboseiterations.state=None: If specified, a previous state of the SAGA solver to continue from (gradient table and averages). Otherwise, the state will be initialized at zeron_ex=None: The total number of examples in the dataloader. If not specified, a single pass will be made over the dataloader to count the total number of examples.n_classes=None: the total number of classes in the dataloader. If not specified, a single pass will be made over the dataloader to count the total number of classes.tol=1e-4: The tolerance level for the stopping criteria of the SAGA solverepsilon=0.001: The regularization path will be calculated at log-spaced intervals betweenmax_lambdaandmax_lambda*epsilon, wheremax_lambdais calculated to be the smallest regularization which results in the all zero solution. The elastic-net paper recommendsepsilon=0.001k=100: The number of steps to take along the regularization pathcheckpoint=None: If specified, save the weights and solver logs for each point of the regularization path within the directorycheckpoint(makes the directory if it does not exist)solve='saga': A string which specifies a particular solver to use (stochastic proximal gradient viasolver=spgis experimental and not recommended)do_zero=True: If true, at the end of the regularization path calculate one more solution corresponding to zero regularization (i.e. fully dense linear model)lr_decay_factor=50: The learning rate of solver will be decayed frommax_lrtomax_lr/lr_decay_factor. Adjust this value to be smaller if progress stalls before reaching an optimal solution, or adjust this value to be larger if the solution path is unstable.metadata=None: a dictionary which contains metadata about the representation which can be used instead ofn_exandn_classes. See below for the assumed structure.val_loader=None: If specified, will calculate statistics (loss and accuracy) and perform model selection based on the given validation settest_loader=None: If specified, will calculate statistics (loss and accuracy) on the given test setlookbehind: The stopping criterion strategy. IfNone, the solver will stop when progress within an interation is less thantol. If specified as an integer, the solver will stop whentolprogress has not been made for more thanlookbehindsteps. The second is more accurate, but will typically take longer.family='multinomial': The distribution family for the GLM. Supported familes aremultinomialandgaussianFor extremely large datasets like ImageNet, it may be desired to avoid multiple initial passes over the dataloader which only calculates dataset statistics. This can be done by supplying themetadatavariable, which is of the following form:metadata = {
    'X' : {
        'mean': 0, 
        'std': 1
    },
    'y' : {
        'mean': 0, 
        'std': 1
    }, 
    'max_reg': {
        'grouped': 0.5, 
        'ungrouped': 0.5
    }
}Any metadata supplied through this variable will not be recomputed. Not all variables need to be specified (i.e. it is possible to supply only the mean and standard deviation, and still perform one pass to calculate the maximum regularization).Additional helper functionsThe package also implements several additional functions which are helpful in order to adapt datasets to the format required by the solver, such as adding example indices and normalizers for the data.Adding indices to datasets and dataloadersIndexedTensorDataset(TensorDataset): 
    def __init__(self, *tensors):A subclass of the PyTorchTensorDatasetwhich returns the tensor indices in additionclass IndexedDataset(Dataset): 
    def __init__(self, ds, sample_weight=None):ADatasetwrapper which takes a PyTorchDatasetwhich returns the dataset indices in additionsample_weight=Nonecan be specified to weight each example differently (e.g. for passing to LIME)add_index_to_dataloader(loader, sample_weight=None):A function which takes a dataloader and returns a new dataloader which returns the dataloader indices in additionsample_weight=Nonecan be specified to weight each example differentlyNormalizing datasetsSometimes a PyTorch dataset or dataloader is not easy to normalize directly. In this case, we can construct a normalizing PyTorch module and pass this into the solver via thepreprocessargument.class NormalizedRepresentation(nn.Module): 
    def __init__(self, loader, model=None, do_tqdm=True, mean=None, std=None, metadata=None, device='cuda'):A module which normalizes inputs by the mean and standard deviation of the given dataloadermodel=NoneIf specified, examples will be passed through the givenmodelbefore calculating the mean and standard deviationdo_tqdm=True: If true, usetqdmprogress barsmean=None: If specified, uses this as the mean instead of calculating the mean from the dataloaderstd=None: If specified, uses this as the standard deviation instead of calculating the standard deviation from the dataloadermetadata=None: If specified, uses dataset statistics from the given dictionarydevice='cuda': The device to store the mean and standard deviation onTestsA number of tests are intests.py, which match the output of the implemented solve to solvers that exist inscikit-learn.