""" transfer_agent.py — Architecture-aware Transfer Learning for LoopUnrollEnv - arch별로 x86 / arm64 등을 선택적으로 지원 - Backbone: 기존 {arch}_base 모델의 일부 레이어를 백본으로 사용 - Adapter: 새 환경(또는 새 CPU)에 맞게 소형 레이어만 재학습 """ import os import glob import sys import argparse import numpy as np import torch import torch.nn as nn import gymnasium as gym from stable_baselines3 import PPO from stable_baselines3.common.env_util import make_vec_env from compiler_env import LoopUnrollEnv # ───────────────────────────────────────────────────────────── # 유틸: 경로 및 기본 설정 # ───────────────────────────────────────────────────────────── PROJECT_ROOT = os.path.expanduser("~/projects/machineai") MODELS_DIR = os.path.join(PROJECT_ROOT, "models") BENCH_DIR = os.path.join(PROJECT_ROOT, "benchmarks") def get_model_paths(arch: str): """ 아키텍처별 기본 모델/전이 모델 경로 생성 - base: models/model_{arch}_base.zip - transfer: models/model_{arch}_transfer.zip """ base = os.path.join(MODELS_DIR, f"model_{arch}_base.zip") transfer = os.path.join(MODELS_DIR, f"model_{arch}_transfer.zip") return base, transfer # ───────────────────────────────────────────────────────────── # Backbone 가중치 추출 # ───────────────────────────────────────────────────────────── def extract_backbone_weights(model_path: str) -> dict: """ 기존 PPO 모델에서 mlp_extractor의 일부 레이어를 백본으로 추출 - 현재는 policy_net의 첫 두 레이어를 백본으로 사용 """ print(f"[Backbone] 로드: {model_path}") model = PPO.load(model_path) state_dict = model.policy.state_dict() backbone = {} for k, v in state_dict.items(): if "mlp_extractor.policy_net.0" in k or "mlp_extractor.policy_net.2" in k: backbone[k] = v.clone() print(f"[Backbone] 추출 레이어:") for k in backbone.keys(): print(f" - {k}") return backbone # ───────────────────────────────────────────────────────────── # Transfer PPO 빌더 # ───────────────────────────────────────────────────────────── def build_transfer_model(env, backbone_weights: dict | None, freeze_backbone: bool = True): """ Backbone 동결 + Adapter 레이어 추가한 PPO 모델 구성 - backbone_weights가 None이면 순수 새 모델로 시작 """ print("[Model] Transfer PPO 생성 중...") model = PPO( policy="MlpPolicy", env=env, learning_rate=1e-4, # 전이학습은 낮은 lr n_steps=256, batch_size=64, n_epochs=10, gamma=0.99, verbose=1, policy_kwargs=dict(net_arch=[64, 64, 32]), # +32 adapter layer ) # 백본 가중치 주입 if backbone_weights is not None: print("[Model] Backbone 가중치 주입...") state_dict = model.policy.state_dict() injected, skipped = 0, 0 for k, v in backbone_weights.items(): if k in state_dict and state_dict[k].shape == v.shape: state_dict[k] = v injected += 1 print(f" ✔ 주입: {k}") else: skipped += 1 print(f" ✗ 스킵: {k} (shape mismatch or not found)") model.policy.load_state_dict(state_dict) print(f"[Model] 주입 완료: {injected}개, 스킵: {skipped}개") else: print("[Model] Backbone 없이 새 모델로 시작") # 백본 동결 if freeze_backbone and backbone_weights is not None: print("[Model] Backbone 파라미터 동결...") for name, param in model.policy.named_parameters(): if "mlp_extractor.policy_net.0" in name or "mlp_extractor.policy_net.2" in name: param.requires_grad = False print(f" 🔒 동결: {name}") trainable = sum(p.numel() for p in model.policy.parameters() if p.requires_grad) total = sum(p.numel() for p in model.policy.parameters()) print(f"\n[Model] 파라미터: {trainable}/{total} 학습가능 ({trainable/total*100:.1f}%)") return model # ───────────────────────────────────────────────────────────── # 메인 전이학습 실행 # ───────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser(description="Architecture-aware transfer learning for LoopUnrollEnv") parser.add_argument("--arch", type=str, default="x86", help="타겟 아키텍처 (예: x86, arm64)") parser.add_argument("--timesteps", type=int, default=2000, help="전이학습 스텝 수") parser.add_argument("--load-base", action="store_true", help="기존 base 모델에서 backbone을 로드할지 여부") parser.add_argument("--base-path", type=str, default="", help="직접 base 모델 경로 지정 (옵션)") parser.add_argument("--out-path", type=str, default="", help="전이 결과 저장 경로 직접 지정 (옵션)") parser.add_argument("--repeat-runs", type=int, default=3, help="실행 시간 측정 반복 횟수") parser.add_argument("--freeze-backbone", action="store_true", help="Backbone 레이어를 동결할지 여부") parser.add_argument("--clang-bin", type=str, default="", help="사용할 clang 바이너리 (비우면 기본값)") parser.add_argument("--opt-bin", type=str, default="", help="사용할 opt 바이너리 (비우면 기본값)") parser.add_argument("--source-files", type=str, nargs="+", default=[], help="학습에 사용할 소스 파일 목록") args = parser.parse_args() arch = args.arch print(f"[Config] arch={arch}") # 경로 설정 os.makedirs(MODELS_DIR, exist_ok=True) default_base, default_transfer = get_model_paths(arch) base_model_path = args.base_path or default_base transfer_model_path = args.out_path or default_transfer print(f"[Config] base_model_path = {base_model_path}") print(f"[Config] transfer_model_path= {transfer_model_path}") # 학습 대상 소스 파일 if args.source_files: source_files = [os.path.abspath(f) for f in args.source_files] else: source_files = sorted(glob.glob(os.path.join(BENCH_DIR, "*.c"))) print(f"[Data] 학습 대상: {source_files}") # Backbone 로드 (옵션) backbone = None if args.load_base: if not os.path.exists(base_model_path): raise FileNotFoundError(f"Base 모델을 찾을 수 없습니다: {base_model_path}") backbone = extract_backbone_weights(base_model_path) else: print("[Backbone] base 모델 로드 생략 (순수 새 모델로 시작)") # Env 생성 함수 def make_env(): return LoopUnrollEnv( source_files=source_files, repeat_runs=args.repeat_runs, arch=arch, clang_bin=args.clang_bin or None, opt_bin=args.opt_bin or None, ) vec_env = make_vec_env(make_env, n_envs=1) # Transfer 모델 빌드 print("\n=== Transfer 모델 빌드 ===") model = build_transfer_model(vec_env, backbone, freeze_backbone=args.freeze_backbone) # 학습 print(f"\n=== Adapter 학습 ({args.timesteps} 스텝) ===") model.learn(total_timesteps=args.timesteps, progress_bar=True) # 저장 model.save(transfer_model_path.replace(".zip", "")) print(f"\n저장 완료: {transfer_model_path}") if __name__ == "__main__": main()