📗 서론
모델을 개발 할 때 모바일 탑재라던지, 사용자의 디스크 용량의 한계 및 처리속도를 위해서 모델의 용량이 매우 작아야 하는 경우가 있다. 이 때 모델 용량은 줄이고 성능은 극대화 할 수 있도록 방안을 찾아보겠다.
1. 경량화 모델 종류
- Yolo 시리즈
- MobileNet
- EfficientNet
- SqueezeNet
- ShuffleNet
2. 모델 선택
YOLO 모델의 경우 v1 ~ v3을 제외한 모델은 GPL 및 AGPL 라이센스를 따르고 있다. 상용화 목적으로 사용 하기 위해서는 라이센스를 구매 하거나 서비스 되고 있는 코드를 오픈소스로 공개 해야 한다.
YOLO의 Backbone을 사용하여 학습한 가중치 모델 또한 라이센스에 적용 되는지에 대한 논쟁이 있지만 가중치 모델 또한 AGPL 라이센스에 포함된다고 명시 되어 있기 때문에 상용화를 목적으로 한다면 이를 고려해야 했다.
나는 우선 상용화를 하겠다는 목적을 가지고 있다고 가정하여 라이센스에 위반되지 않는 MobileNet을 우선 사용 해보도록 하겠다.
3. MobileNet v4
가장 최근에 나온 모델인 MobileNet v4를 이용하여 학습을 진행 해 보겠다.
1) MobileNet v4 구현 코드
이 알고리즘을 구현한 코드는 jaiwei98/MobileNetV4-pytorch를 참고하여 구현 하였다.
2) 학습 코드 구현
위에서 구현된 코드를 보면 MNV4ConvSmall_BLOCK_SPECS, MobileNetV4ConvMedium, MobileNetV4ConvLarge 등 가장 적은 수의 파라미터를 갖는 small모델부터 많은 파라미터를 갖는 Large모델을 볼 수 있다. 이 모델 정의를 바탕으로 학습을 진행 해보겠다.
- 학습 파라미터
우선 아래와 같이 학습 시 파라미터를 받아올 수 있도록 코드를 작성 해보았다.
parser = argparse.ArgumentParser(description="Train or Fine-Tune MobileNetV4")
parser.add_argument("--train_data", type=str, required=True, help="Path to dataset root directory")
parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
parser.add_argument("--lr", type=float, default=0.001, help="Learning rate")
parser.add_argument("--img_size", type=int, default=224, help="Input image size (default: 224x224)")
parser.add_argument("--model_type", type=str, default="MobileNetV4ConvSmall_v2",
choices=["MobileNetV4ConvSmall", "MobileNetV4ConvMedium", "MobileNetV4ConvLarge"],
help="MobileNetV4 model type")
parser.add_argument("--save_path", type=str, default="./checkpoints/mobilenetv4",
help="Path to save the trained model")
parser.add_argument("--finetune", action="store_true", help="Enable fine-tuning (use pretrained model)")
parser.add_argument("--pretrained", type=str, default=None, help="Path to pretrained model (.pth) for fine-tuning")
parser.add_argument("--freeze", "--freeze_layers", action="store_true", help="Freeze lower layers during fine-tuning")
args = parser.parse_args()--train_data : 학습 데이터 경로
--epochs : 학습 반복 횟수
--batch_size : 배치 사이즈
--lr : 학습률
--img_size : Input 이미지의 사이즈
--model_type : small, medium, large 선택
--save_path : 모델 저장경로
--finetune : 모델 finetuning 여부
--pretrained : 모델 finetuning시 사전 학습 모델 경로 지정
--freeze : 모델 finetuning시 Layer를 freeze할지 안할지 선택-
데이터 로드 및 학습
-
우선 데이터를 불러온다 get_dataloaders 함수를 통해 데이터 전처리 및 train, valid 데이터로 나눈다.
-
이후 model을 정의 한다. MobileNetv4에는 총 6개의 레이어로 구성 되어 있고 convolution layer와 MobileNet에서 주로 사용하는 UIB 레이어가 존재 한다.
model = MobileNetV4(model=args.model_type)학습 시 파라미터로 small, medium, large중에 선택하면 된다.
# 데이터 로더 설정
def get_dataloaders(data_dir, batch_size, img_size, valid_ratio=0.2):
transform = transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
full_dataset = datasets.ImageFolder(root=data_dir, transform=transform)
total_size = len(full_dataset)
valid_size = int(total_size * valid_ratio)
train_size = total_size - valid_size
train_dataset, valid_dataset = random_split(full_dataset, [train_size, valid_size])
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
return train_loader, valid_loader, full_dataset.classes
train_loader, valid_loader, class_names = get_dataloaders(args.train_data, args.batch_size, args.img_size)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 모델 생성
model = MobileNetV4(model=args.model_type)
# Fine-Tuning 모드일 경우
if args.finetune:
if args.pretrained is None:
raise ValueError("⚠️ Fine-Tuning을 위해 --pretrained 모델 경로를 지정해야 합니다!")
print(f"🔄 Loading pretrained model from {args.pretrained} for fine-tuning...")
model.load_state_dict(torch.load(args.pretrained, map_location=device)) # 가중치 불러오기
# 특정 레이어 Freeze 설정 (Feature Extraction 모드)
freeze_layers(model, args.freeze_layers)
# 옵티마이저 (Freeze 여부에 따라 학습할 레이어만 업데이트)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr)
else:
print(" Training from scratch...")
optimizer = optim.Adam(model.parameters(), lr=args.lr)
criterion = nn.CrossEntropyLoss()
# 학습 시작 (Train or Fine-Tune)
train(model, train_loader, valid_loader, criterion, optimizer, device, args.epochs, args.save_path)-
finetuning시 동작
-
freeze layer를 통해 첫 3개의 layer를 고정하여 기존 모델의 성질을 보존한다.
-
만약 새로 추가된 데이터가 많거나 새로운 특징이 더욱 학습이 많이 진행되어야 한다면 layer를 고정시키지 않고 학습을 돌릴 수 도 있다.
def freeze_layers(model, freeze=True):
"""
특정 레이어를 Freeze(고정)하는 함수
:param model: MobileNetV4 모델
:param freeze: True면 일부 레이어 고정, False면 전체 학습 가능
"""
if freeze:
print("🔒 Freezing lower layers... (Feature Extraction Mode)")
for param in model.conv0.parameters(): # 초기 Conv Layer 고정
param.requires_grad = False
for param in model.layer1.parameters(): # Layer1 고정
param.requires_grad = False
for param in model.layer2.parameters(): # Layer2 고정
param.requires_grad = False
else:
print("🔓 All layers are trainable!")이후 train 함수를 통해 학습을 진행 한다. 이 부분은 정해진 방법없기 때문에 본인 취향에 맞게 코드를 작성 하면 된다. 나는 model.train()의 metrics 정보들을 tqdm 함수를 이용하여 진행률을 확인 하는 방식으로 진행 하였다. 추후에는 MLflow를 통해 가시화 시킬 예정이다.
3. MobileNetv4 조금 더 경량화
위의 방식으로 학습을 진행 해 보았는데 모델 사이즈가 10MB정도 나왔다. 모바일에 탑재 되어야 한다고 생각하면 좀 더 줄이면 좋을 것 같아서 모델 구조를 더 바꿔 보았다.
기존 model 구조
MNV4ConvSmall_BLOCK_SPECS = {
"conv0": {
"block_name": "convbn",
"num_blocks": 1,
"block_specs": [
[3, 32, 3, 2]
]
},
"layer1": {
"block_name": "convbn",
"num_blocks": 2,
"block_specs": [
[32, 32, 3, 2],
[32, 32, 1, 1]
]
},
"layer2": {
"block_name": "convbn",
"num_blocks": 2,
"block_specs": [
[32, 96, 3, 2],
[96, 64, 1, 1]
]
},
"layer3": {
"block_name": "uib",
"num_blocks": 6,
"block_specs": [
[64, 96, 5, 5, True, 2, 3],
[96, 96, 0, 3, True, 1, 2],
[96, 96, 0, 3, True, 1, 2],
[96, 96, 0, 3, True, 1, 2],
[96, 96, 0, 3, True, 1, 2],
[96, 96, 3, 0, True, 1, 4],
]
},
"layer4": {
"block_name": "uib",
"num_blocks": 6,
"block_specs": [
[96, 128, 3, 3, True, 2, 6],
[128, 128, 5, 5, True, 1, 4],
[128, 128, 0, 5, True, 1, 4],
[128, 128, 0, 5, True, 1, 3],
[128, 128, 0, 3, True, 1, 4],
[128, 128, 0, 3, True, 1, 4],
]
},
"layer5": {
"block_name": "convbn",
"num_blocks": 2,
"block_specs": [
[128, 960, 1, 1],
[960, 1280, 1, 1]
]
}
}수정된 모델 구조
MNV4ConvSmall_BLOCK_SPECS_v2 = {
"conv0": {
"block_name": "convbn",
"num_blocks": 1,
"block_specs": [
[3, 16, 3, 2] # 32 → 16
]
},
"layer1": {
"block_name": "convbn",
"num_blocks": 2,
"block_specs": [
[16, 16, 3, 2], # 32 → 16
[16, 16, 1, 1]
]
},
"layer2": {
"block_name": "depthwise_separable",
"num_blocks": 2,
"block_specs": [
[16, 48, 3, 2], # 96 → 48
[48, 32, 1, 1] # 64 → 32
]
},
"layer3": {
"block_name": "uib",
"num_blocks": 6,
"block_specs": [
[32, 64, 5, 5, True, 2, 3], # 96 → 64
[64, 64, 0, 3, True, 1, 2],
[64, 64, 0, 3, True, 1, 2],
[64, 64, 0, 3, True, 1, 2],
[64, 64, 0, 3, True, 1, 2],
[64, 64, 3, 0, True, 1, 4],
]
},
"layer4": {
"block_name": "uib",
"num_blocks": 6,
"block_specs": [
[64, 96, 3, 3, True, 2, 6], # 128 → 96
[96, 96, 5, 5, True, 1, 4],
[96, 96, 0, 5, True, 1, 4],
[96, 96, 0, 5, True, 1, 3],
[96, 96, 0, 3, True, 1, 4],
[96, 96, 0, 3, True, 1, 4],
]
},
"layer5": {
"block_name": "convbn",
"num_blocks": 2,
"block_specs": [
[96, 768, 1, 1], # 128 → 96, 960 → 768
[768, 960, 1, 1] # 960 → 768, 1280 → 960
]
}
}-
기존에 conv layer를 사용하는 대신 depthwise_separable를 사용하면 연산량과 메모리 사용량을 줄일 수 있고 모바일이나 임베디드 환경에서 유리한 학습이 가능하다.
-
단, 복잡한 패턴에서는 학습 능력이 떨어질 수 있고 ResNet계열에서는 depthwise_separable보다는 conv 레이어가 적합할 수 있다.
-
depthwise_separable 구현 코드는 아래와 같다
class DepthwiseSeparableConv(nn.Module):
def __init__(self, inp, oup, kernel_size=3, stride=1, padding=1, bias=False):
super().__init__()
self.depthwise = nn.Conv2d(inp, inp, kernel_size, stride, padding, groups=inp, bias=bias)
self.pointwise = nn.Conv2d(inp, oup, 1, bias=bias)
self.bn = nn.BatchNorm2d(oup)
self.act = nn.ReLU(inplace=True)
def forward(self, x):
x = self.depthwise(x)
x = self.pointwise(x)
x = self.bn(x)
return self.act(x)- 이 코드를 MobileNet의 모델 정의 코드 중 build_blocks 함수에 해당 layer를 추가할 수 있도록 수정 해주면 된다.
def build_blocks(layer_spec):
if not layer_spec.get('block_name'):
return nn.Sequential()
block_names = layer_spec['block_name']
layers = nn.Sequential()
if block_names == "convbn":
schema_ = ['inp', 'oup', 'kernel_size', 'stride']
for i in range(layer_spec['num_blocks']):
args = dict(zip(schema_, layer_spec['block_specs'][i]))
layers.add_module(f"convbn_{i}", conv_2d(**args))
elif block_names == "depthwise_separable":
schema_ = ['inp', 'oup', 'kernel_size', 'stride']
for i in range(layer_spec['num_blocks']):
args = dict(zip(schema_, layer_spec['block_specs'][i]))
layers.add_module(f"convbn_{i}", DepthwiseSeparableConv(**args))
...
...
...
return layers이렇게 수정 하여 모델 학습을 진행 해 보니 약 5.8MB의 모델이 나왔다.
내가 준비한 테스트 셋 약 1200개를 각 모델로 테스트 해본 결과
10MB 모델 인식률은 94%,
5.8MB의 인식률은 91.5% 정도로 2.5% 정도의 정확도 하락은 있었지만 경량화 된 모델의 장점이 있기 때문에 상황에 맞게 조절 해가며 학습을 진행 하면 될 것 같다.
