예제 소스
실시간 XGBoost와 DQN 결합 구조: 거래 AI를 검증하는 법 · 파일 59개
PngToStatus.py
"""
차트 PNG → 상태(잠재 표현) → 액션 로짓을 만드는 정책 네트워크.
원래 ``PngToStatus`` (gym 환경 호환용) 외에, 학습 가능한 ``PngToStatusPolicy``
와 PNG 바이트 입출력 헬퍼를 함께 제공합니다.
핵심:
- **`PngToStatusPolicy`**: 캡쳐 차트 PNG → CNN 임베딩 + (보유여부·미실현손익) 사이드 특징
→ 매도/매수/유지 3-class 로짓. REINFORCE/모방학습 등에 그대로 사용 가능.
- **`png_bytes_to_tensor`**: 메모리 PNG 바이트를 그대로 모델 입력 텐서로 변환
(디스크 저장 없이 빠르게 학습/추론 가능).
"""
from __future__ import annotations
import io
import logging
from pathlib import Path
from typing import Optional
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as T
from PIL import Image
logger = logging.getLogger(__name__)
def _device_default() -> torch.device:
"""기본 device.
NOTE: Apple MPS 는 ``AdaptiveAvgPool2d`` 에서 input/output 크기가
나누어떨어지지 않으면 RuntimeError 가 발생한다(PyTorch issue #96056).
이 모델은 매우 작아 MPS 의 가속 이득이 작고, CPU↔MPS 전송 오버헤드가
오히려 더 크므로 CPU 를 기본으로 사용한다. CUDA 만 자동 선택.
"""
if torch.cuda.is_available():
return torch.device("cuda")
return torch.device("cpu")
# ---------------------------------------------------------------------------
# 학습 가능한 정책 (메인)
# ---------------------------------------------------------------------------
DEFAULT_INPUT_H = 96
DEFAULT_INPUT_W = 192
N_ACTIONS = 3 # 0=매도, 1=매수, 2=유지 (main_llm_db_test_ollama._opinion_to_signal 와 동일)
SIDE_DIM = 3 # in_position(0/1), unrealized_pnl, cash_ratio
class PngToStatusPolicy(nn.Module):
"""차트 PNG(메모리 바이트) → 매도/매수/유지 로짓 정책.
forward 는 항상 **로짓**을 반환합니다(샘플링/argmax 는 호출측에서).
이전 ``PngToStatus`` 와 달리, 학습이 가능합니다.
"""
def __init__(
self,
h: int = DEFAULT_INPUT_H,
w: int = DEFAULT_INPUT_W,
n_actions: int = N_ACTIONS,
side_dim: int = SIDE_DIM,
hidden: int = 256,
in_channels: int = 3,
) -> None:
super().__init__()
self.h = int(h)
self.w = int(w)
self.n_actions = int(n_actions)
self.side_dim = int(side_dim)
# ── feature backbone ── (AdaptiveAvgPool 은 MPS 호환을 위해 forward 에서 처리)
self.features = nn.Sequential(
nn.Conv2d(in_channels, 32, kernel_size=5, stride=2, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=5, stride=2, padding=2),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(2),
nn.Conv2d(64, 96, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(96),
nn.ReLU(inplace=True),
)
# 입력 해상도와 무관하게 같은 차원이 되도록 글로벌 평균 풀링 사용.
# (MPS 에서도 output_size=(1,1) 은 항상 안전.)
self._pool_out = (1, 1)
self._feat_channels = 96
self.head = nn.Sequential(
nn.Flatten(),
nn.Linear(self._feat_channels * self._pool_out[0] * self._pool_out[1] + self.side_dim, hidden),
nn.ReLU(inplace=True),
nn.Dropout(0.3),
nn.Linear(hidden, self.n_actions),
)
self._transform = T.Compose(
[
T.Resize((self.h, self.w)),
T.ToTensor(), # 0..1, [C,H,W]
]
)
@property
def device(self) -> torch.device:
return next(self.parameters()).device
def forward(self, x: torch.Tensor, side: torch.Tensor) -> torch.Tensor:
"""x: [B,3,H,W] side: [B, side_dim] → logits: [B, n_actions]"""
feat = self.features(x)
# MPS 의 adaptive_avg_pool2d 는 input/output 비-나누어떨어지면 미지원.
# output (1,1) 자체는 안전하지만, 혹시 다른 출력 크기로 바꿀 때를 대비해
# 안전하게 CPU 폴백 경로를 둔다.
if feat.device.type == "mps" and (
feat.shape[-2] % self._pool_out[0] != 0
or feat.shape[-1] % self._pool_out[1] != 0
):
pooled = F.adaptive_avg_pool2d(feat.to("cpu"), self._pool_out).to(feat.device)
else:
pooled = F.adaptive_avg_pool2d(feat, self._pool_out)
flat = pooled.flatten(1)
if side.dim() == 1:
side = side.unsqueeze(0)
merged = torch.cat([flat, side.to(flat.dtype)], dim=1)
return self.head(merged)
def png_bytes_to_tensor(self, png_bytes: bytes) -> torch.Tensor:
"""메모리 PNG → 모델 입력 텐서 [1,3,H,W] (현재 device)."""
im = Image.open(io.BytesIO(png_bytes)).convert("RGB")
x = self._transform(im).unsqueeze(0)
return x.to(self.device)
def make_side_tensor(
self,
*,
in_position: bool,
unrealized_pnl: float,
cash_ratio: float,
) -> torch.Tensor:
"""포트폴리오 보조 특징 텐서 [1, side_dim]."""
vec = [
1.0 if in_position else 0.0,
float(np.clip(unrealized_pnl, -1.0, 1.0)),
float(np.clip(cash_ratio, 0.0, 1.0)),
]
return torch.tensor([vec], dtype=torch.float32, device=self.device)
# ---------------------------------------------------------------------------
# 저장 / 로드
# ---------------------------------------------------------------------------
def save_png_status_policy(model: PngToStatusPolicy, path: str | Path) -> None:
p = Path(path)
p.parent.mkdir(parents=True, exist_ok=True)
torch.save(model.state_dict(), p)
logger.info("PngToStatusPolicy 저장: %s", p)
def load_png_status_policy(
path: str | Path,
*,
h: int = DEFAULT_INPUT_H,
w: int = DEFAULT_INPUT_W,
n_actions: int = N_ACTIONS,
side_dim: int = SIDE_DIM,
hidden: int = 256,
device: Optional[torch.device] = None,
) -> PngToStatusPolicy:
dev = device or _device_default()
model = PngToStatusPolicy(
h=h, w=w, n_actions=n_actions, side_dim=side_dim, hidden=hidden
).to(dev)
p = Path(path)
if p.exists():
try:
state = torch.load(p, map_location=dev)
model.load_state_dict(state)
logger.info("PngToStatusPolicy 로드: %s", p)
except Exception as e: # 차원 변경 등은 새로 시작
logger.warning("PngToStatusPolicy 로드 실패(%s) → 새 가중치", e)
else:
logger.info("PngToStatusPolicy 가중치 없음 → 신규 모델")
return model
# ---------------------------------------------------------------------------
# 구버전 PngToStatus / Env (호환용 — 사용처 없으면 제거 가능)
# ---------------------------------------------------------------------------
def save_png_to_status_model(model: nn.Module, filepath: str | Path = "pt/png_to_status.pt") -> None:
Path(filepath).parent.mkdir(parents=True, exist_ok=True)
torch.save(model.state_dict(), filepath)
def load_png_to_status_model(model: nn.Module, filepath: str | Path = "pt/png_to_status.pt") -> None:
model.load_state_dict(torch.load(filepath))
class PngToStatus(nn.Module):
"""구버전 호환용. 학습용으로는 ``PngToStatusPolicy`` 사용을 권장합니다."""
def __init__(self, h: int, w: int, qsize: int) -> None:
super().__init__()
hidden = 128
self.hight = h
self.width = w
self.conv1 = nn.Conv2d(4, h * w, kernel_size=5, stride=2)
self.bn1 = nn.BatchNorm2d(h * w)
self.conv2 = nn.Conv2d(h * w, qsize, kernel_size=5, stride=2)
self.bn2 = nn.BatchNorm2d(qsize)
self.conv3 = nn.Conv2d(qsize, qsize, kernel_size=3, stride=2, padding=1)
self.bn3 = nn.BatchNorm2d(qsize)
self.flatten = nn.Flatten()
self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
linear_input_size = 3200
self.fc = nn.Linear(linear_input_size, hidden)
self.dropout = nn.Dropout(0.5)
self.head = nn.Linear(hidden, qsize)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.pool(F.relu(self.bn1(self.conv1(x))))
x = self.pool(F.relu(self.bn2(self.conv2(x))))
x = self.flatten(x)
x = torch.sigmoid(x)
return torch.argmax(x, dim=1)