[Typing] Literal, tensortype ( torch typing )

Literal

python 3.8부터 추가된 기능으로 Literal typing을 통해서 특정 변수가 가질 수 있는 변수의 범위를 제한할 수 있다. 즉, 기존의 static typing의 장점을 가져오기 위해 사용하던 python의 type hinting을 넘어서 dynamic한 파이썬의 특징에 맞게 다양한 타입에 해당되더라도 사용자의 목적성에 맞게 원하는 변수들로 값을 제한할 수 있는 것이다.

물론, 파이썬이기에 Literal을 통해서 제한해둔 값의 범위를 넘어간 값을 입력하는 것도 가능하다. 즉 실제 런타임 상에서 에러가 발생하지는 않는다. 따라서, Literal로 범위를 제한해두고, 함수 내에서 조건문을 통해서 exception handling을 해주는 것이 일반적이다. 

from typing import Literal

def func(a : Literal[1,'3']):
    if a in Literal[1,'3']:
        return a
    else:
        raise ValueError('invalid input')

 

torchtyping

torch tensor의 shape, dtype, names 등에 대해서 type annotation을 제공하는 오픈소스이다. 종종 볼륨이 큰 오픈소스들에서 보이기에 정리해두고자 한다.

크게 다음과 같은 annoation들을 제공한다.

• shape: size, number of dimensions;
• dtype (float, integer, etc.);
• layout (dense, sparse);
• names of dimensions as per named tensors;
• arbitrary number of batch dimensions with ...;
• ...plus anything else you like, as torchtyping is highly extensible.

torch를 사용하다보면 종종 shape의 불일치로 인해서 squeeze 혹은 unsqueeze를 빼먹어 불필요한 오류가 발생하는 상황을 겪게 된다. 협업 그리고 스스로도 이런 상황을 피하기 위해서는 torch tensor에 대해서 typing을 해두는 것이 좋다고 생각한다. 다음과 같이 TensorType이라는 type을 사용해 tensor가 가져야하는 shape에 대해서 일종의 annotation을 제공할 수 있다. x라는 tensor는 ( B,C )와 같은 형식일 것이다. 실제로 많은 저장소들에서 이를 주석으로 처리했었다. 더 바람직한 방법은 다음과 같이 typing을 하는 것이라고 생각한다.

def batch_outer_product(x:   TensorType["batch", "x_channels"],
                        y:   TensorType["batch", "y_channels"]
                        ) -> TensorType["batch", "x_channels", "y_channels"]:

    return x.unsqueeze(-1) * y.unsqueeze(-2)

실제로 torchtyping에 적혀있는 문구다 무척 유쾌하다.

with programmatic checking that the shape (dtype, ...) specification is met.
Bye-bye bugs! Say hello to enforced, clear documentation of your code.

If (like me) you find yourself littering your code with comments like # x has shape (batch, hidden_state) or statements like assert x.shape == y.shape , just to keep track of what shape everything is, then this is for you.

실제 사용은 다음과 같다. typeguard를 미리 설치해둔다면 runtime에서 type checking까지 진행할 수 있다. 보다 restric한 typing이 적용되게 할 수 있다는 것이다. 하단의 예제를 보면 TensorType에 "batch"라는 shape annotation을 했다. 그렇다면, batch에 해당되는 임의의 값의 shape은 x와 y에서 consistent해야한다. 따라서 가장 마지막 line처럼 shape을 다르게 설정하면 오류가 발생하는 것이다. 

from torch import rand
from torchtyping import TensorType, patch_typeguard
from typeguard import typechecked

patch_typeguard()  # use before @typechecked

@typechecked
def func(x: TensorType["batch"],
         y: TensorType["batch"]) -> TensorType["batch"]:
    return x + y

func(rand(3), rand(3))  # works
func(rand(3), rand(1))
# TypeError: Dimension 'batch' of inconsistent size. Got both 1 and 3.

필자의 사용을 위해 몇가지를 더 정리하도록 하겠다.

shape

• int를 사용하게 되면 명확히 해당 사이즈로 되어야한다. -1을 사용하게 되면 어떤 사이즈도 가능하도록 설정하는 것이다.
• str을 사용하게 되면 ( 위 예시 ) runtime에서 넘어가는 argument의 size가 이 name으로 bound되도록 하는 것이다. 따라서 해당 name ( 위에서 'batch')으로 타이핑되어 있는 값들에 대해서 모두 size가 consistent한지를 체킹한다.
• ... : any sizes에 대해서 임의의 수의 dimension을 처리할 때 사용한다. ( ex - [8,3,256,256] => [8, ...] 이런식으로 타이핑할 수 있겠다. )

dtype

• int, bool, float : float의 경우 torch.get_default_dtype()에서 설정된 float값을 사용한다. ( 일반적으로 float32 )

 

 

 

 

Reference

https://yogyui.tistory.com/entry/PythontypingLiteral-%EB%B3%80%EC%88%98-%ED%98%95type-%EB%B0%8F-%EA%B0%92-%EC%A0%9C%ED%95%9C%ED%95%98%EA%B8%B0

Python::typing::Literal - 변수값 범위 명시

https://github.com/patrick-kidger/torchtyping

GitHub - patrick-kidger/torchtyping: Type annotations and dynamic checking for a tensor's shape, dtype, names, etc.