def get_predictions(model, dataloaders, device="cpu", as_dict=False, per_neuron=True, test_data=True, **kwargs): predictions = {} with eval_state(model) if not isinstance( model, types.FunctionType) else contextlib.nullcontext(): for k, v in dataloaders.items(): if test_data: _, output = model_predictions_repeats(dataloader=v, model=model, data_key=k, device=device) else: _, output = model_predictions(dataloader=v, model=model, data_key=k, device=device) predictions[k] = output.T if not as_dict: predictions = [v for v in predictions.values()] return predictions
def get_targets(model, dataloaders, device="cpu", as_dict=True, per_neuron=True, test_data=True, **kwargs): responses = {} with eval_state(model) if not isinstance( model, types.FunctionType) else contextlib.nullcontext(): for k, v in dataloaders.items(): if test_data: targets, _ = model_predictions_repeats(dataloader=v, model=model, data_key=k, device=device) targets_per_neuron = [] for i in range(targets[0].shape[1]): neuronal_responses = [] for repeats in targets: neuronal_responses.append(repeats[:, i]) targets_per_neuron.append(neuronal_responses) responses[k] = targets_per_neuron else: targets, _ = model_predictions(dataloader=v, model=model, data_key=k, device=device) responses[k] = targets.T if not as_dict: responses = [v for v in responses.values()] return responses
def get_poisson_loss(model, dataloaders, device="cpu", as_dict=False, avg=False, per_neuron=True, eps=1e-12): poisson_loss = {} with eval_state(model) if not isinstance( model, types.FunctionType) else contextlib.nullcontext(): for k, v in dataloaders.items(): target, output = model_predictions(dataloader=v, model=model, data_key=k, device=device) loss = output - target * np.log(output + eps) poisson_loss[k] = np.mean(loss, axis=0) if avg else np.sum(loss, axis=0) if as_dict: return poisson_loss else: if per_neuron: return np.hstack([v for v in poisson_loss.values()]) else: return (np.mean(np.hstack([v for v in poisson_loss.values()])) if avg else np.sum(np.hstack([v for v in poisson_loss.values()])))
def get_correlations(model, dataloaders, device="cpu", as_dict=False, per_neuron=True, **kwargs): correlations = {} with eval_state(model) if not isinstance( model, types.FunctionType) else contextlib.nullcontext(): for k, v in dataloaders.items(): target, output = model_predictions(dataloader=v, model=model, data_key=k, device=device) correlations[k] = corr(target, output, axis=0) if np.any(np.isnan(correlations[k])): warnings.warn("{}% NaNs , NaNs will be set to Zero.".format( np.isnan(correlations[k]).mean() * 100)) correlations[k][np.isnan(correlations[k])] = 0 if not as_dict: correlations = (np.hstack([v for v in correlations.values()]) if per_neuron else np.mean( np.hstack([v for v in correlations.values()]))) return correlations
def get_module_output(model, input_shape, use_cuda=True): """ Returns the output shape of the model when fed in an array of `input_shape`. Note that a zero array of shape `input_shape` is fed into the model and the shape of the output of the model is returned. Args: model (nn.Module): PyTorch module for which to compute the output shape input_shape (tuple): Shape specification for the input array into the model use_cuda (bool, optional): If True, model will be evaluated on CUDA if available. Othewrise model evaluation will take place on CPU. Defaults to True. Returns: tuple: output shape of the model """ # infer the original device initial_device = next(iter(model.parameters())).device device = "cuda" if torch.cuda.is_available() and use_cuda else "cpu" with eval_state(model): with torch.no_grad(): input = torch.zeros(1, *input_shape[1:], device=device) output = model.to(device)(input) model.to(initial_device) return output.shape
def get_FEV(dataloaders, model, device='cpu', as_dict=False): FEV = {} with eval_state(model) if not isinstance( model, types.FunctionType) else contextlib.nullcontext(): for k, v in dataloaders.items(): repeated_inputs, repeated_outputs = get_repeats(v) FEV[k] = compute_FEV( repeated_inputs=repeated_inputs, repeated_outputs=repeated_outputs, model=model, device=device, data_key=k, ) return FEV if as_dict else np.hstack([v for v in FEV.values()])
def model_predictions_repeats(dataloader, model, data_key, device='cpu', broadcast_to_target=False): """ Computes model predictions for dataloader that yields batches with identical inputs along the first dimension Unique inputs will be forwarded only once through the model Returns: target: ground truth, i.e. neuronal firing rates of the neurons as a list: [num_images][num_reaps, num_neurons] output: responses as predicted by the network for the unique images. If broadcast_to_target, returns repeated outputs of shape [num_images][num_reaps, num_neurons] else (default) returns unique outputs of shape [num_images, num_neurons] """ output = torch.empty(0) target = [] # Get unique images and concatenate targets: unique_images = torch.empty(0) for images, responses in dataloader: if len(images.shape) == 5: images = images.squeeze(dim=0) responses = responses.squeeze(dim=0) assert torch.all(torch.eq( images[-1, ], images[0, ], )), "All images in the batch should be equal" unique_images = torch.cat((unique_images, images[0:1, ]), dim=0) target.append(responses.detach().cpu().numpy()) # Forward unique images once: with eval_state(model) if not isinstance( model, types.FunctionType) else contextlib.nullcontext(): with device_state(model, device): output = model(unique_images.to(device), data_key=data_key).detach().cpu() output = output.numpy() if broadcast_to_target: output = [ np.broadcast_to(x, target[idx].shape) for idx, x in enumerate(output) ] return target, output
def get_module_output(model, input_shape): """ Gets the output dimensions of the convolutional core by passing an input image through all convolutional layers :param core: convolutional core of the DNN, which final dimensions need to be passed on to the readout layer :param input_shape: the dimensions of the input :return: output dimensions of the core """ initial_device = "cuda" if next(iter( model.parameters())).is_cuda else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu" with eval_state(model): with torch.no_grad(): input = torch.zeros(1, *input_shape[1:]).to(device) output = model.to(device)(input) model.to(initial_device) return output.shape
def get_FEV(model, dataloaders, device="cpu", as_dict=False, per_neuron=True, threshold=None): """ Computes the fraction of explainable variance explained (FEVe) per Neuron, given a model and a dictionary of dataloaders. The dataloaders will have to return batches of identical images, with the corresponing neuronal responses. Args: model (object): PyTorch module dataloaders (dict): Dictionary of dataloaders, with keys corresponding to "data_keys" in the model device (str): 'cuda' or 'gpu as_dict (bool): Returns the scores as a dictionary ('data_keys': values) if set to True. per_neuron (bool): Returns the grand average if set to True. threshold (float): for the avg feve, excludes neurons with a explainable variance below threshold Returns: FEV (dict, or np.array, or float): Fraction of explainable varianced explained. Per Neuron or as grand average. """ dataloaders = dataloaders["test"] if "test" in dataloaders else dataloaders FEV = {} with eval_state(model) if not isinstance( model, types.FunctionType) else contextlib.nullcontext(): for data_key, dataloader in dataloaders.items(): targets, outputs = model_predictions_repeats( model=model, dataloader=dataloader, data_key=data_key, device=device, broadcast_to_target=True) if threshold is None: FEV[data_key] = compute_FEV(targets=targets, outputs=outputs) else: fev, feve = compute_FEV(targets=targets, outputs=outputs, return_exp_var=True) FEV[data_key] = feve[fev > threshold] if not as_dict: FEV = np.hstack([v for v in FEV.values()]) if per_neuron else np.mean( np.hstack([v for v in FEV.values()])) return FEV