Design patterns for machine learning in a nutshell (1): Creational Design Pattern

Creational design patterns and their motivation

Creational design pattern: Instance 생성과 관련된 패턴. 어떤 Instance를 효율적으로 생성하고 관리할 수 있게 한다.

 

Creational pattern은 위 처럼 instance 생성이라는 목적으로 만들어진 것들이다. 그렇다면, 모델 학습 파이프라인 등을 구성하는데에는 어떠한 예시가 있을 수 있을까? Instance 생성이 효과적이고, 직관적으로 되는 케이스가 어떤 것들이 있을까?

HuggingFace의 AutoModel 같은 경우를 생각해 보자. 이는 (End-user의 입장에서) 모델 instance를 생성하는 데에 굉장히 편리한 역할을 한다: AutoModel.from_pretrained. 혹은, 어떤 object를 생성하는데에 이것저것 복잡한 로직이 들어간다고 생각해 보자. torch.nn.Module instance를 만들기 전에 configuration을 파싱 하고, accelerate, deep-speed를 쓸 건지, Single-gpu를 쓸 건지 따위의 것들이 들어가거나, Dataloader instance를 만들기 전에 필요한 몇 가지 로직이 있는 경우가 있겠다. 혹은 새로운 데이터로더 로직을 작성하려고 하는데, 기존에 존재하는 코드를 수정하면서 작성하고 싶지 않은 상황이 있을 수 있겠다.

이 포스팅 에서는, ML 학습을 주로 예시로 들어서 몇 가지 creational pattern들을 설명한다. 전통적으로 사용되어 온 패턴들을 소개하고, 중간중간 Huggingface 등에서 보이는 같은 목적을 가진 modern-pattern 들도 다룬다.

Patterns

Factory Patterns

여러 클래스들의 instance를 생성하는 목적을 가진 패턴. 특정 학습방법론을 실험하는데, 여러 architecture 에서 모두 잘 working 하는지를 보고 싶다고 생각해 보자. 다른 모든 logic에 대한 코드는 같고, model을 불러오는 부분만 다르다고 할 때, naive 하게는 분기로써 아래와 같이 구현할 수 있겠다.

def create_model(model_type):
    if model_type == "CNN":
        return nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 10)
        )
    elif model_type == "RNN":
        return nn.Sequential(
            nn.RNN(input_size=10, hidden_size=20, num_layers=2, batch_first=True)
        )
    elif model_type == "Transformer":
        return nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6)
    else:
        raise ValueError("Unknown model type")

 
Factory Pattern은 위처럼 학습을 하는 Main code에서, 이런 분기를 타는 create_model을 만들지 말고, ModelFactory 라는 별도의 클래스가 모델 인스턴스 생성을 하도록 만들어보자. 아래와 같이, 여러 종류의 instance를 만들 수 있는 하나의 class를 가지는것을 Simple factory pattern 이라한다.

class ModelFactory:
    def create_model(self, model_type: str, some_complex_config: SomeConfig):
        if model_type == "CNN":
        	...
            return CNNModel()
        elif model_type == "RNN":
            ...
            return RNNModel()
        elif model_type == "Transformer":
        	...
            return TransformerModel()
        else:
            raise ValueError("Unknown model type")

class CNNModel:
    def train(self):
        print("Training CNN model...")

class RNNModel:
    def train(self):
        print("Training RNN model...")

class TransformerModel:
    def train(self):
        print("Training Transformer model...")

factory = ModelFactory()
model = factory.create_model("CNN", some_complex_configuration)
model.train()

 
하지만, 이는 여전히 팩토리 클래스 내부적으로는 -분기-를 기반으로 하고있다. 더 나은 해결책으로는, 아래처럼 모델마다 모델 생성 팩토리를 만들 수 있겠다. 이러한 패턴은 Factory method pattern 이라 불리운다.

from abc import ABC, abstractmethod
import torch.nn as nn

class ModelFactory(ABC):
    @abstractmethod
    def create_model(self):
        pass

class CNNFactory(ModelFactory):
    def create_model(self):
        return nn.Sequential(...)

class RNNFactory(ModelFactory):
    def create_model(self):
        return nn.Sequential(nn.RNN(input_size=10, hidden_size=20, num_layers=2, batch_first=True))

class TransformerFactory(ModelFactory):
    def create_model(self):
        return nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6)

def get_model(factory: ModelFactory):
    model = factory.create_model()
    return model

cnn_factory = CNNFactory()
cnn_model = get_model(cnn_factory)
print("CNN Model:", cnn_model)

rnn_factory = RNNFactory()
rnn_model = get_model(rnn_factory)
print("RNN Model:", rnn_model)

transformer_factory = TransformerFactory()
transformer_model = get_model(transformer_factory)
print("Transformer Model:", transformer_model)

 
이제, 새로운 모델이 생길때마다 분기를 만드는 짓을 하지 않아도 된다. 그런데, 얼핏 보면, 새로운 모델을 추가할 때마다 새로운 무언가를 추가해야 하는 건 팩토리를 만드나 안 만드나 똑같고 (해당하는 팩토리 클래스를 만들어야 하니까), 복잡성만 증가한 것이 아닌가 생각되기도 한다.
 
실제로는, 이렇게 분리하게 되면 아래와 같은 몇 장점이 있다.

1. 새로운 모델을 추가 할 때, 완전히 별도의 Factory class에서 진행하게 되면 기존 코드를 잘 못 건드릴 일이 없다. 즉, 안정성이 증가한다.
2. Open/Closed Pricinciple (OCP) 준수: 즉, 기존 코드를 전혀 수정하지 않고 새로운 기능을 확장할 수 있다. 내가 구현하려는 새로운 모델의 클래스에 대해서만 구현에 신경을 쓰면 되는 것이다. 기존의 Interface (Base Class or Abstract class of the factory)를 수정하지 않아도 된다. (혹은, 않아야 한다!)
3. 동일 인터페이스를 따를 수 밖에 없기 때문에 (ModelFactory), 코드에 일관된 패턴을 강제할 수 있다.

working with the classmethod: AutoClass

전통적인 의미의 Factory 패턴은 아니지만, 비슷한 목적을 가지지만 Python의 유연성을 활용한 디자인을 한번 알아보자. Python의 classmethod를 이용해서, Factory패턴과 유사한 목적-인스턴스 생성-을 위해 마치 Huggingface의 AutoClass처럼 사용할 수도 있다:

class AutoModel:
    @classmethod
    def create_model(cls, model_type):
        if model_type == "CNN":
            return CNNModel()
        elif model_type == "RNN":
            return RNNModel()
        elif model_type == "Transformer":
            return TransformerModel()
        else:
            raise ValueError("Unknown model type")

...

model = AutoModel.create_model("CNN")
model.train()  # Output: Training CNN model...

 
여전히 분기가 거슬린다고하면, HF에서 처럼 Model-registery를 활용해서 해당 분기를 없애고, 모델을 추가할 때 아예 별도의 Factory를 만들 필요도 없게 만들 수도 있다! 각 모델의 팩토리를 만드는 것이 아니고, 그냥 Class decorator를 추가하기만 하면 된다.

class AutoModel:
    _registry = {}

    @classmethod
    def register_model(cls, model_type: str):
        def inner_wrapper(wrapped_class):
            cls._registry[model_type] = wrapped_class
            return wrapped_class
        return inner_wrapper

    @classmethod
    def from_type(cls, model_type: str, some_complex_config):
        if model_type not in cls._registry:
            raise ValueError(f"Unknown model type: {model_type}")
        return cls._registry[model_type](some_complex_config)

@AutoModel.register_model("CNN")
class CNNModel:
    def __init__(self, config):
        print("Initializing CNN model with config:", config)
    def train(self):
        print("Training CNN model...")

@AutoModel.register_model("RNN")
class RNNModel:
    def __init__(self, config):
        print("Initializing RNN model with config:", config)
    def train(self):
        print("Training RNN model...")

@AutoModel.register_model("Transformer")
class TransformerModel:
    def __init__(self, config):
        print("Initializing Transformer model with config:", config)
    def train(self):
        print("Training Transformer model...")

model = AutoModel.from_type("CNN", some_complex_configuration)
model.train()

 
이는 원래의 Factory pattern의 장점인 OCP, 객체의 생성과 클래스 기능의 분리가 없다. 하지만, 코드가 간결해지고 '라이브러리 구조' 에 적합하다. 즉, End-user는 아묻따 AutoSomething.give_my_instance 만 하면 되는 것이다.

Abstract Factory Pattern

이름에서 알 수 있듯, Factory Pattern처럼 "instance를 생성하는 기능"을 하는 클래스를 가지는 패턴이다. 다만, Abstract Factory Class에서는 instance들을 생성한다. ML 학습을 예시로 들면, 어떤 Class가 데이터로더, 모델, 학습루프 (trainer)와 관련된 instance를 모조리 생성하는 것이다. 

다양한 실험 환경 구성을 한다고 해보자. Abstract Factory Pattern을 사용하면 아마 아래와 같이 구현할 수 있을 것이다.

from abc import ABC, abstractmethod
import torch
import torch.nn as nn
import torch.optim as optim

# Abstract Factory
class ExperimentFactory(ABC):
    @abstractmethod
    def create_model(self):
        pass
    
    @abstractmethod
    def create_data_loader(self):
        pass
    
    @abstractmethod
    def create_trainer(self):
        pass

class CNNExperimentFactory(ExperimentFactory):
    def create_model(self):
        return nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 10)
        )

    def create_data_loader(self):
        return torch.utils.data.DataLoader(
            torch.randn(100, 3, 32, 32), batch_size=32, shuffle=True)

    def create_trainer(self):
        return Trainer(self.create_model(), self.create_data_loader())

class RNNExperimentFactory(ExperimentFactory):
    def create_model(self):
        return nn.RNN(input_size=10, hidden_size=20, num_layers=2, batch_first=True)
    
    def create_data_loader(self):
        return torch.utils.data.DataLoader(
            torch.randn(100, 10, 10), batch_size=32, shuffle=True)

    def create_trainer(self):
        return Trainer(self.create_model(), self.create_data_loader())

class Trainer:
    def __init__(self, model, data_loader):
        self.model = model
        self.data_loader = data_loader
        self.optimizer = optim.SGD(model.parameters(), lr=0.01)
    
    def train_loop(self):
        print(f"Starting training with {self.model.__class__.__name__}")
        for epoch in range(N): 
            for data in self.data_loader:
                outputs = self.model(data)
                ...
                loss.backward()
                self.optimizer.step()
                self.optimizer.zero_grad()
        print("Training completed.")

cnn_factory = CNNExperimentFactory()
cnn_trainer = cnn_factory.create_trainer()
cnn_trainer.train_loop()

print("\n---\n")

rnn_factory = RNNExperimentFactory()
rnn_trainer = rnn_factory.create_trainer()
rnn_trainer.train_loop()

 
위에서는 create_data_loader가 단순히 DataLoader class가 들어갔지만, 실제로는 Dataloader의 Factory가 들어갈 수 있겠다.
 
이렇게 하면, Trainer 같은 공통 로직을 재사용할 수 있고, 확장성이 좋아진다. 즉, 새로운 모델과 데이터로 실험을 한다고 할 때, 기존 코드를 손대지 않고 (즉, 기존 코드에 문제가 발생할 일은 없다.) 새로운 ExperimentFactory만 추가하면 된다.

Singleton Pattern

특정 클래스가 단 하나의 instance만 가지도록 강제하는 패턴. 무언가 공유되어야 하는 리소스 (log, gpu resource) 등을 다룰 때 유용하다. 예를 들어, torch로 학습을 할 때는 cpu, cuda:0 등의 device 설정을 model.to(device)로 해 줄 때가 종종 있고, 이런 Device는 모든 code에 걸쳐 동일해야 한다 (하나의 Device setting을 가져야 한다).
Accelerate 등을 사용하지 않는다면 말이다. 아래의 예시처럼 구현할 수 있다:

import torch

class DeviceManager:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            print("Initializing Device Manager...")
            cls._instance = super(DeviceManager, cls).__new__(cls)
            cls._instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        return cls._instance

    def get_device(self):
        return self.device

device_manager1 = DeviceManager()
print("Device:", device_manager1.get_device())  # cuda or cpu

device_manager2 = DeviceManager()
print("Device:", device_manager2.get_device()) 

print(device_manager1 is device_manager2)  # True, since two different device_managers share the instance

 

Surpassing global variable

이 때, Singleton Pattern이 Global variable과 유사하다고 생각될 수 있다.

 

하지만 Singleton pattern은 global variable로 하나의 값을 관리하는 것에 비해 장점이 몇가지 있다:

1. Singleton object는 "필요할 때만 초기화" 할 수 있고 (gpu 여부를 항상 체크하지 않아도 될 수 있다)
2. 관련 설정을 더 추가하고 싶을 때 용이하며 ( 특정 파라미터 수 미만의 모델은 CPU로 돌리고 싶거나, config를 받아서 testing일 때는 CPU로 돌리거나
3. global variable과 달리 mocking이 용이해서 테스트가 더 쉬워진다.

Builder Pattern

Tensor Flow나 torch의 Sequential을 생각하면 쉽다. 복잡한 instance를 여러 단계로 구성할 때 사용된다. 가장 간단한 예시는 아래와 같다.

import torch.nn as nn

class CNNModelBuilder:
    def __init__(self):
        self.layers = []

    def add_conv_layer(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        self.layers.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding))
        return self

    def add_relu(self):
        self.layers.append(nn.ReLU())
        return self

    def add_max_pool(self, kernel_size, stride):
        self.layers.append(nn.MaxPool2d(kernel_size, stride))
        return self

    def add_fc_layer(self, in_features, out_features):
        self.layers.append(nn.Linear(in_features, out_features))
        return self

    def build(self):
        return nn.Sequential(*self.layers)

builder = CNNModelBuilder()
model = (
    builder.add_conv_layer(3, 16, 3)
           .add_relu()
           .add_max_pool(2, 2)
           .add_conv_layer(16, 32, 3)
           .add_relu()
           .add_max_pool(2, 2)
           .add_fc_layer(32 * 6 * 6, 10)
           .build()
)

print(model)

 
실제로는, 2020년도 초반에 ML research를 하면 항상 볼 수 있던 ResNet code block에서 이러한 로직을 많이 사용한다. 다만, 실제로 ML프로젝트에 class로써 builder pattern 적용된 케이스는 많지 않은 듯하다.