[PyTorch] Weight clipping

이전 글에서 PyTorch를 이용한 gradient clipping을 소개했는데, 

gradient가 아니라 weight 값 자체를 일정 범위 안으로 제한해야 하는 경우가 있다.

 

이 경우 Weight clipping을 수행해주는 class를 정의하여 사용할 수 있다.

 

출처: https://discuss.pytorch.org/t/set-constraints-on-parameters-or-layers/23620

Set constraints on parameters or layers

 

 

Weight clipper 추가하기

다음과 같이 Weight clipper을 정의해줄 수 있다. torch.clamp method를 이용하면 tensor의 값을 해당 범위 안으로 clip해줄 수 있다.

class WeightClipper(object):
    def __call__(self, module, param, clip_min, clip_max):
    	if hasattr(module, param):
	    self.clip(module, param, clip_min, clip_max)
        
    def clip(self, module, param, clip_min, clip_max):
        p = getattr(module, param).data
        p = p.clamp(clip_max, clip_max)
        getattr(module, param).data = p

 

사용법은 아래와 같다.

model = Net()
weightclipper = WeightClipper()

#after each backward operation
loss.backward()
optimizer.step()
model.apply(weightclipper)

 

 

---

 

나의 경우는 model 내의 어떤 parameter이 특정 범위 안으로 제한되도록 위와 같이 weight clipper을 정의했지만,

만약 conv2d 등의 모든 module들의 weight를 [-1, 1] 범위로 clip하고 싶다면 다음과 같이 만들 수도 있다.

class WeightClipper(object):
    def __call__(self, module, param):
    	if hasattr(module, param):
	    self.clip(module, param)
        
    def clip(self, module, param):
        p = getattr(module, param).data
        p = p.clamp(-1, 1)
        getattr(module, param).data = p