Skip to content

util

model.util

get_last_conv_channels

get_last_conv_channels(model: Module) -> int

WARNING: This will return the number of output channels of the last conv layer defined in model, and it could differ from the actual output channels of the model.

Parameters:

  • model (Module) –

    torch.nn.Module who will return its last conv layer channels

Returns:

  • int ( int ) –

    last conv layer channels

Source code in SaigeToolkit/model/util.py
def get_last_conv_channels(model: torch.nn.Module) -> int:
    """WARNING: This will return the number of output channels of the last conv layer defined in `model`,
    and it could differ from the actual output channels of the model.

    Args:
        model (torch.nn.Module): torch.nn.Module who will return its last conv layer channels

    Returns:
        int: last conv layer channels
    """

    layer_type_list = [torch.nn.Conv2d, torch.nn.ConvTranspose2d]
    return [module for module in model.modules() if type(module) in layer_type_list][-1].out_channels

extend_state_dict_input_channel

extend_state_dict_input_channel(state_dict: Mapping[str, Any], input_weight_key: str, input_conv_layer: Conv2d) -> None

(Multipage) 3채널 이상인 이미지를 사용하기 위해 필요한 기능이며, state_dict의 input conv 채널이 네트워크의 input conv 채널보다 작은 경우 해당 weight의 채널을 늘려줍니다. 현재 네트워크가 가진 input conv 웨이트에서 앞 3 채널을 state_dict의 input conv 웨이트로 치환하는 방식을 사용합니다.

Parameters:

  • state_dict (Mapping[str, Any]) –

    로드하려는 weight

  • input_weight_key (str) –

    input conv weight의 이름

  • input_conv_layer (Conv2d) –

    input conv layer

Source code in SaigeToolkit/model/util.py
def extend_state_dict_input_channel(
    state_dict: Mapping[str, Any],
    input_weight_key: str,
    input_conv_layer: nn.Conv2d,
) -> None:
    """(Multipage) 3채널 이상인 이미지를 사용하기 위해 필요한 기능이며,
    state_dict의 input conv 채널이 네트워크의 input conv 채널보다 작은 경우 해당 weight의 채널을 늘려줍니다.
    현재 네트워크가 가진 input conv 웨이트에서 앞 3 채널을 state_dict의 input conv 웨이트로 치환하는 방식을 사용합니다.

    Args:
        state_dict (Mapping[str, Any]): 로드하려는 weight
        input_weight_key (str): input conv weight의 이름
        input_conv_layer (nn.Conv2d): input conv layer
    """
    state_dict_in_channels = state_dict[input_weight_key].shape[1]
    if state_dict_in_channels != input_conv_layer.in_channels:
        state_dict_in_channels = min(state_dict_in_channels, input_conv_layer.in_channels)
        original_input_conv_weight = input_conv_layer.state_dict()["weight"]
        original_input_conv_weight[:, :state_dict_in_channels] = state_dict[input_weight_key][
            :, :state_dict_in_channels
        ]
        state_dict[input_weight_key] = original_input_conv_weight