| import torch
|
| from torch import nn
|
| from collections import OrderedDict
|
|
|
|
|
| def to_cuda(tensors):
|
| """Transfer tensor, dict or list of tensors to GPU.
|
|
|
| Args:
|
| tensors (:class:`torch.Tensor`, list or dict): May be a single, a
|
| list or a dictionary of tensors.
|
|
|
| Returns:
|
| :class:`torch.Tensor`:
|
| Same as input but transferred to cuda. Goes through lists and dicts
|
| and transfers the torch.Tensor to cuda. Leaves the rest untouched.
|
| """
|
| if isinstance(tensors, torch.Tensor):
|
| return tensors.cuda()
|
| if isinstance(tensors, list):
|
| return [to_cuda(tens) for tens in tensors]
|
| if isinstance(tensors, dict):
|
| for key in tensors.keys():
|
| tensors[key] = to_cuda(tensors[key])
|
| return tensors
|
| raise TypeError(
|
| "tensors must be a tensor or a list or dict of tensors. "
|
| " Got tensors of type {}".format(type(tensors))
|
| )
|
|
|
|
|
| def tensors_to_device(tensors, device):
|
| """Transfer tensor, dict or list of tensors to device.
|
|
|
| Args:
|
| tensors (:class:`torch.Tensor`): May be a single, a list or a
|
| dictionary of tensors.
|
| device (:class: `torch.device`): the device where to place the tensors.
|
|
|
| Returns:
|
| Union [:class:`torch.Tensor`, list, tuple, dict]:
|
| Same as input but transferred to device.
|
| Goes through lists and dicts and transfers the torch.Tensor to
|
| device. Leaves the rest untouched.
|
| """
|
| if isinstance(tensors, torch.Tensor):
|
| return tensors.to(device)
|
| elif isinstance(tensors, (list, tuple)):
|
| return [tensors_to_device(tens, device) for tens in tensors]
|
| elif isinstance(tensors, dict):
|
| for key in tensors.keys():
|
| tensors[key] = tensors_to_device(tensors[key], device)
|
| return tensors
|
| else:
|
| return tensors
|
|
|
|
|
| def pad_x_to_y(x, y, axis=-1):
|
| """Pad first argument to have same size as second argument
|
|
|
| Args:
|
| x (torch.Tensor): Tensor to be padded.
|
| y (torch.Tensor): Tensor to pad x to.
|
| axis (int): Axis to pad on.
|
|
|
| Returns:
|
| torch.Tensor, x padded to match y's shape.
|
| """
|
| if axis != -1:
|
| raise NotImplementedError
|
| inp_len = y.size(axis)
|
| output_len = x.size(axis)
|
| return nn.functional.pad(x, [0, inp_len - output_len])
|
|
|
|
|
| def load_state_dict_in(state_dict, model):
|
| """Strictly loads state_dict in model, or the next submodel.
|
| Useful to load standalone model after training it with System.
|
|
|
| Args:
|
| state_dict (OrderedDict): the state_dict to load.
|
| model (torch.nn.Module): the model to load it into
|
|
|
| Returns:
|
| torch.nn.Module: model with loaded weights.
|
|
|
| # .. note:: Keys in a state_dict look like object1.object2.layer_name.weight.etc
|
| We first try to load the model in the classic way.
|
| If this fail we removes the first left part of the key to obtain
|
| object2.layer_name.weight.etc.
|
| Blindly loading with strictly=False should be done with some logging
|
| of the missing keys in the state_dict and the model.
|
|
|
| """
|
| try:
|
|
|
|
|
| model.load_state_dict(state_dict, strict=True)
|
| except RuntimeError:
|
|
|
|
|
|
|
|
|
|
|
| new_state_dict = OrderedDict()
|
| for k, v in state_dict.items():
|
| new_k = k[k.find(".") + 1 :]
|
| new_state_dict[new_k] = v
|
| model.load_state_dict(new_state_dict, strict=True)
|
| return model
|
|
|
|
|
| def are_models_equal(model1, model2):
|
| """Check for weights equality between models.
|
|
|
| Args:
|
| model1 (nn.Module): model instance to be compared.
|
| model2 (nn.Module): second model instance to be compared.
|
|
|
| Returns:
|
| bool: Whether all model weights are equal.
|
| """
|
| for p1, p2 in zip(model1.parameters(), model2.parameters()):
|
| if p1.data.ne(p2.data).sum() > 0:
|
| return False
|
| return True
|
|
|
