pyiqa.archs.arch_util ===================== .. py:module:: pyiqa.archs.arch_util Module Contents --------------- .. py:function:: dist_to_mos(dist_score: torch.Tensor) -> torch.Tensor Convert distribution prediction to MOS score. For datasets with detailed score labels, such as AVA. :param dist_score: (*, C), C is the class number. :type dist_score: torch.Tensor :returns: (*, 1) MOS score. :rtype: torch.Tensor .. py:function:: random_crop(input_list, crop_size, crop_num) Randomly crops the input tensor(s) to the specified size and number of crops. :param input_list: List of input tensors or a single input tensor. :type input_list: list or torch.Tensor :param crop_size: Size of the crop. If an int is provided, a square crop of that size is used. If a tuple is provided, a crop of that size is used. :type crop_size: int or tuple :param crop_num: Number of crops to generate. :type crop_num: int :returns: If a single input tensor is provided, a tensor of cropped images is returned. If a list of input tensors is provided, a list of tensors of cropped images is returned. :rtype: torch.Tensor or list of torch.Tensor .. py:function:: uniform_crop(input_list, crop_size, crop_num) Crop the input_list of tensors into multiple crops with uniform steps according to input size and crop_num. :param input_list: List of input tensors or a single input tensor. :type input_list: list or torch.Tensor :param crop_size: Size of the crops. If int, the same size will be used for height and width. If tuple, should be (height, width). :type crop_size: int or tuple :param crop_num: Number of crops to generate. :type crop_num: int :returns: Cropped tensors. If input_list is a list, the output will be a list of cropped tensors. If input_list is a single tensor, the output will be a single tensor. :rtype: torch.Tensor or list of torch.Tensor .. py:function:: clip_preprocess_tensor(x: torch.Tensor, model) Clip preprocess function with tensor input. NOTE: Results are slightly different with original preprocess function with PIL image input, because of differences in resize function. :param x: Input tensor. :type x: torch.Tensor :param model: Model with visual input resolution. :returns: Preprocessed tensor. :rtype: torch.Tensor .. py:function:: clean_state_dict(state_dict) Clean checkpoint by removing .module prefix from state dict if it exists from parallel training. :param state_dict: State dictionary from a model checkpoint. :type state_dict: dict :returns: Cleaned state dictionary. :rtype: dict .. py:function:: get_url_from_name(name: str, store_base: str = 'hugging_face', base_url: str = None) -> str Get the URL for a given file name from a specified storage base. :param name: The name of the file. :type name: str :param store_base: The storage base to use. Options are "hugging_face" or "github". Default is "hugging_face". :type store_base: str, optional :param base_url: Base URL to use if provided. :type base_url: str, optional :returns: The URL of the file. :rtype: str .. py:function:: load_pretrained_network(net: torch.nn.Module, model_path: str, strict: bool = True, weight_keys: str = None) -> None Load a pretrained network from a given model path. :param net: The network to load the weights into. :type net: torch.nn.Module :param model_path: Path to the model weights file. Can be a URL or a local file path. :type model_path: str :param strict: Whether to strictly enforce that the keys in state_dict match the keys returned by net's state_dict(). Default is True. :type strict: bool, optional :param weight_keys: Specific key to extract from the state_dict. Default is None. :type weight_keys: str, optional :returns: None .. py:data:: to_1tuple .. py:data:: to_2tuple .. py:data:: to_3tuple .. py:data:: to_4tuple .. py:data:: to_ntuple .. py:function:: default_init_weights(module_list, scale=1, bias_fill=0, **kwargs) Initializes the weights of the given module(s) using Kaiming Normal initialization. :param module_list: List of modules or a single module to initialize. :type module_list: list or nn.Module :param scale: Scaling factor for the weights. Default is 1. :type scale: float, optional :param bias_fill: Value to fill the biases with. Default is 0. :type bias_fill: float, optional :param \*\*kwargs: Additional arguments for the Kaiming Normal initialization. :returns: None .. rubric:: Example >>> import torch.nn as nn >>> from arch_util import default_init_weights >>> model = nn.Sequential( >>> nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1), >>> nn.ReLU(), >>> nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), >>> nn.ReLU(), >>> nn.Linear(64 * 32 * 32, 10) >>> ) >>> default_init_weights(model, scale=0.1, bias_fill=0.01)