fedsim.fl.utils — fedsim 0.1 documentation

# Source code for fedsim.fl.utils

    from functools import partial
    from inspect import signature
    import torch
    from torch.nn.utils import clip_grad_norm_
    from torch.nn.utils.convert_parameters import _check_param_device

    def get_metric_scores(metric_fn_dict, y_true, y_pred):
        answer = {}
        if metric_fn_dict is None:
            return answer
        for name, fn in metric_fn_dict.items():
            args = dict()
            if 'y_true' in signature(fn).parameters:
                args['y_true'] = y_true
            elif 'target' in signature(fn).parameters:
                args['target'] = y_true
            else:
                raise NotImplementedError
            if 'y_pred' in signature(fn).parameters:
                args['y_pred'] = y_pred
            elif 'input' in signature(fn).parameters:
                args['input'] = y_pred
            else:
                raise NotImplementedError
            answer[name] = fn(**args)
        return answer

    def default_closure(x,
                        y,
                        model,
                        loss_fn,
                        optimizer,
                        metric_fn_dict,
                        max_grad_norm=1000,
                        link_fn=partial(torch.argmax, dim=1),
                        device='cpu',
                        transform_grads=None,
                        transform_y=None,
                        **kwargs):
        if transform_y is not None:
            y = transform_y(y)
        y_true = y.tolist()
        x = x.to(device)
        y = y.reshape(-1).long()
        y = y.to(device)
        model.train()
        outputs = model(x)
        loss = loss_fn(outputs, y)
        if loss.isnan() or loss.isinf():
            return loss
        # backpropagation
        loss.backward()
        if transform_grads is not None:
            transform_grads(model)
        # Clip gradients
        clip_grad_norm_(parameters=model.parameters(), max_norm=max_grad_norm)
        # optimize
        optimizer.step()
        optimizer.zero_grad()
        y_pred = link_fn(outputs).tolist()
        metrics = get_metric_scores(metric_fn_dict, y_true, y_pred)
        return loss, metrics

    def vector_to_parameters_like(vec, parameters_like):
        r"""Convert one vector to new parameters like the ones provided

        Args:
            vec (Tensor): a single vector represents the parameters of a model.
            parameters (Iterable[Tensor]): an iterator of Tensors that are the
                parameters of a model. This is only used to get the sizes. New
                parametere are defined.
        """
        # Ensure vec of type Tensor
        if not isinstance(vec, torch.Tensor):
            raise TypeError('expected torch.Tensor, but got: {}'.format(
                torch.typename(vec)))
        # Flag for the device where the parameter is located
        param_device = None

        # Pointer for slicing the vector for each parameter
        pointer = 0
        new_params = []
        for param in parameters_like:
            # Ensure the parameters are located in the same device
            param_device = _check_param_device(param, param_device)

            # The length of the parameter
            num_param = param.numel()
            # Slice the vector, reshape it, and replace the old data of the
            # parameter
            new_params.append(vec[pointer:pointer + num_param].view_as(param).data)

            # Increment the pointer
            pointer += num_param
        return new_params

    class ModelReconstructor(torch.nn.Module):

        def __init__(self,
                     feature_extractor,
                     classifier,
                     connection_fn=None) -> None:
            super(ModelReconstructor, self).__init__()
            self.feature_extractor = feature_extractor
            self.classifier = classifier
            self.connection_fn = connection_fn

        def forward(self, input):
            features = self.feature_extractor(input)
            if self.connection_fn is not None:
                features = self.connection_fn(features)
            return self.classifier(features)