계속: Trax에서 사용자 정의 활성화 함수 사용

개시하다

이 문장은 다음 문장의 후속이다.
https://zenn.dev/catminusminus/articles/cf0c81857810f3
마지막으로 Google에서 공개한 Deep Learning 프로그램 라이브러리 중 하나인 Trax에는 독자적인 활성화 함수를 사용하는 방법이 적혀 있습니다.
하지만 그때는 순함수만 처리했다.
이번에 우리는 더욱 복잡한 행위 활성화 함수를 실현할 것이다.

learnable의 경우: Parame tricRelu

Trax에서는 ParametricRelu가 구현되고 있습니다.
그렇다면 단독으로 실시할 의미는 없지만 실제로 Trax의 ParametricRelu는

@assert_shape('...->...')  # The output and input shapes are the same.
def ParametricRelu(a=1.):
  r"""Returns a layer that computes a ReLU function with the given slope.
  .. math::
      f(x) = \left\{ \begin{array}{cl}
          0  & \text{if}\ x \leq 0, \\
          ax & \text{otherwise}.
      \end{array} \right.
  Args:
    a: Slope of line for positive inputs.
  """
  return Fn('ParametricRelu', lambda x: jnp.maximum(a * x, jnp.zeros_like(x)))

및 정의됨.
물론 이것은 Learnable가 아니다.
그렇다면 Learnable의 Parame tricRelu는 어떻게 설치하는 것이 좋을까요?
여기 참고로 다음입니다.

class ThresholdedLinearUnit(base.Layer):
  """Thresholded Linear Unit, c.f. https://arxiv.org/pdf/1911.09737.pdf ."""

  def init_weights_and_state(self, input_signature):
    """Initializes this layer's single weight to zero."""
    del input_signature
    self.weights = jnp.zeros((), dtype=jnp.float32)

  def forward(self, inputs):
    """Executes this layer as part of a forward pass through the model.
    Args:
      inputs: Tensor.
    Returns:
      Tensor of same shape and dtype as the input.
    """
    threshold = self.weights
    return jnp.maximum(inputs, threshold)

즉, 일반적으로 Layer로 구현하면 된다.
그래서 다음과 같은 내용을 실시해 봤다.

@assert_shape("...->...")
class ParametricRelu(Layer):
    def __init__(self, num_parameters: int = 1, a: float = 0.25):
        super().__init__()
        self._a = a
        self._num_parameters = num_parameters

    def init_weights_and_state(self, input_signature):
        del input_signature
        self.weights = self._a * jnp.ones((self._num_parameters,), dtype=jnp.float32)

    def forward(self, inputs):
        return jnp.where(inputs <= 0, self.weights * inputs, inputs)

여기num_parameters는 PyTorch의 인터페이스에 적합하다.
내가 대체로 본 바와 같이 중요한 곳이 하나 있다.
밑줄에는 self._a 및 self._num_parameters처럼 변수가 있습니다.
이것은 필수적이다.
밑줄self.a 등을 넣지 않으면 동작할 수 없습니다.self.weights 등 특별한 것만 허용됩니다.

훈련 여부에 따라 행동이 달라지는 상황: RandomizedRelu

훈련할 때와 그 외에 행동을 바꾸고 싶은 경우가 있다.
전형적인 것은 Droput이다.
그런 상황에서 어떻게 하면 좋을까.
Droput은 forward 함수에서 self._mode가'train'분기 중인지 여부에 따라 달라진다.

  def forward(self, x):
    """Executes this layer as part of a forward pass through the model.
    Args:
      x: Tensor of activations.
    Returns:
      Tensor of same shape and dtype as the input.
    """
    if self._mode != 'train':
      return x

그렇군요. RandomizedRelu를 설치해 보고 싶습니다.

@assert_shape("...->...")
class RandomizedRelu(Layer):
    def __init__(self, lower=0.125, upper=0.3333333333333333, mode="train"):
        super().__init__()
        self._lower = lower
        self._upper = upper
        self._mode = mode

    def forward(self, inputs):
        if self._mode == "train":
            a = fastmath.random.uniform(
                self.rng, dtype=jnp.float32, minval=self._lower, maxval=self._upper
            )
        else:
            a = (self._lower + self._upper) / 2
        return jnp.where(inputs <= 0, a * inputs, inputs)

이대로 만사 해결의 방식으로 보도를 끝내려고 했는데 문제가 있었다.
이러한 실현에서 이를 사용하는 쪽은mode의 전환을 해야 한다.
그리고trainerlib.py내자동 전환 모드.

    model_train = model(mode='train')
    model_predict_eval = model(mode='eval')

그러나, 이trainerlib.py에 있는 것은 구 API입니다.
새 API는 training입니다.py쪽에 있지만 모드의 전환점을 찾을 수 없습니다.
가까운 곳이 있지만 기록된metrics의 전환에만 사용되며 TrainTask와 EvalTask로 모델을 자동으로 전환하는 행위는 불가능할 것 같습니다.
아직 새로운 쪽이'under development'이기 때문에 머지않아 실현될 수 있을 것이다.
단지 내가 잘못 보았을 뿐이니, 그때는 반드시 나에게 메시지를 남겨 주세요.

끝말

따라서 트랙스에서 조금 복잡한 활성화 함수를 구현하는 방법을 소개했다.