Upload transfer_agent.py with huggingface_hub

transfer_agent.py

@@ -0,0 +1,195 @@
+"""
+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()