1{
2  "op": "replace",
3  "filename": "qlib/custom_model.py",
4  "start_line": 16,
5  "end_line": 103,
6  "content": "import copy, torch.optim as optim\n\n\nclass GRUAttn(nn.Module):\n    \"\"\"GRU with temporal attention pooling for stock features.\"\"\"\n    def __init__(self, d_feat=6, hid=64, n_layers=2, drop=0.0):\n        super().__init__()\n        self.d_feat = d_feat\n        self.gru = nn.GRU(d_feat, hid, n_layers, batch_first=True, dropout=drop)\n        self.att = nn.Linear(hid, 1)\n        self.fc = nn.Linear(hid, 1)\n    def forward(self, x):\n        x = x.reshape(len(x), self.d_feat, -1).permute(0, 2, 1)  # [N,T,F]\n        h, _ = self.gru(x)                                         # [N,T,H]\n        w = torch.softmax(self.att(h), dim=1)                      # [N,T,1]\n        return self.fc((h * w).sum(1)).squeeze()                    # [N]\n\n\nclass CustomModel(Model):\n    \"\"\"GRU-Attention with IC-MSE joint loss for stock prediction.\n    Novel: differentiable Pearson-IC loss term optimizes ranking quality.\"\"\"\n    def __init__(self):\n        super().__init__()\n        self.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n        self.net = GRUAttn(d_feat=6, hid=64, n_layers=2, drop=0.0).to(self.device)\n        self.opt = optim.Adam(self.net.parameters(), lr=1e-3, weight_decay=1e-5)\n        self.n_epochs, self.bs, self.early_stop = 200, 2048, 20\n        self.lam = 1.0  # IC loss weight\n        self.fitted = False\n    @staticmethod\n    def _ic(p, t):\n        \"\"\"Differentiable Pearson correlation (IC).\"\"\"\n        pc, tc = p - p.mean(), t - t.mean()\n        return (pc * tc).sum() / (pc.norm() * tc.norm() + 1e-8)\n    def _loss(self, pred, label):\n        m = ~torch.isnan(label)\n        p, t = pred[m], label[m]\n        if len(p) < 2: return (p - t).pow(2).mean()\n        return (p - t).pow(2).mean() - self.lam * self._ic(p, t)\n    def _run(self, x, y, train):\n        self.net.train(train)\n        idx = np.arange(len(x))\n        if train: np.random.shuffle(idx)\n        total, cnt = 0.0, 0\n        cm = torch.enable_grad() if train else torch.no_grad()\n        with cm:\n            for s in range(0, len(idx), self.bs):\n                b = idx[s:s + self.bs]\n                xb = torch.FloatTensor(x[b]).to(self.device)\n                yb = torch.FloatTensor(y[b]).to(self.device)\n                out = self.net(xb)\n                loss = self._loss(out, yb)\n                if train:\n                    self.opt.zero_grad()\n                    loss.backward()\n                    torch.nn.utils.clip_grad_value_(self.net.parameters(), 3.0)\n                    self.opt.step()\n                total += loss.item() * len(b)\n                cnt += len(b)\n        return total / cnt\n    def fit(self, dataset):\n        df_t = dataset.prepare(\"train\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L)\n        df_v = dataset.prepare(\"valid\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L)\n        xt = np.nan_to_num(df_t[\"feature\"].values)\n        yt = np.nan_to_num(np.squeeze(df_t[\"label\"].values))\n        xv = np.nan_to_num(df_v[\"feature\"].values)\n        yv = np.nan_to_num(np.squeeze(df_v[\"label\"].values))\n        best_loss, best_ep = np.inf, 0\n        best_w = copy.deepcopy(self.net.state_dict())\n        for ep in range(self.n_epochs):\n            self._run(xt, yt, train=True)\n            vl = self._run(xv, yv, train=False)\n            if vl < best_loss:\n                best_loss, best_ep = vl, ep\n                best_w = copy.deepcopy(self.net.state_dict())\n            if ep - best_ep >= self.early_stop: break\n        self.net.load_state_dict(best_w)\n        self.fitted = True\n    def predict(self, dataset, segment=\"test\"):\n        if not self.fitted: raise ValueError(\"not fitted\")\n        x_df = dataset.prepare(segment, col_set=\"feature\", data_key=DataHandlerLP.DK_I)\n        x = np.nan_to_num(x_df.values)\n        self.net.eval()\n        parts = []\n        with torch.no_grad():\n            for s in range(0, len(x), self.bs):\n                xb = torch.FloatTensor(x[s:s + self.bs]).to(self.device)\n                parts.append(self.net(xb).reshape(-1).cpu().numpy())\n        return pd.Series(np.concatenate(parts), index=x_df.index)"
7}

1{
2  "op": "replace",
3  "filename": "qlib/custom_model.py",
4  "start_line": 16,
5  "end_line": 104,
6  "content": "import copy, torch.optim as optim\n\n\nclass GRUAttn(nn.Module):\n    \"\"\"GRU with temporal attention and time-step masking.\"\"\"\n    def __init__(self, d_feat=6, hid=64, n_layers=2):\n        super().__init__()\n        self.d_feat = d_feat\n        self.gru = nn.GRU(d_feat, hid, n_layers, batch_first=True)\n        self.att = nn.Linear(hid, 1)\n        self.fc = nn.Linear(hid, 1)\n    def forward(self, x, mask_p=0.0):\n        x = x.reshape(len(x), self.d_feat, -1).permute(0, 2, 1)\n        if mask_p > 0 and self.training:\n            m = (torch.rand(x.shape[0], x.shape[1], 1, device=x.device) > mask_p).float()\n            x = x * m / (1 - mask_p)\n        h, _ = self.gru(x)\n        w = torch.softmax(self.att(h), dim=1)\n        return self.fc((h * w).sum(1)).squeeze()\n\n\nclass CustomModel(Model):\n    \"\"\"GRU-Attention with cross-sectional IC loss and temporal masking.\n    Novel: (1) date-aligned batching for true cross-sectional IC,\n    (2) temporal masking for robust feature learning.\"\"\"\n    def __init__(self):\n        super().__init__()\n        self.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n        self.net = GRUAttn(d_feat=6, hid=64, n_layers=2).to(self.device)\n        self.opt = optim.Adam(self.net.parameters(), lr=1e-3, weight_decay=5e-5)\n        self.n_epochs, self.early_stop, self.lam, self.acc = 200, 20, 2.0, 4\n        self.mask_p, self.fitted = 0.1, False\n    @staticmethod\n    def _ic(p, t):\n        pc, tc = p - p.mean(), t - t.mean()\n        return (pc * tc).sum() / (pc.norm() * tc.norm() + 1e-8)\n    def _group(self, df):\n        gs = []\n        for _, g in df.groupby(level=0):\n            x, y = np.nan_to_num(g[\"feature\"].values), np.nan_to_num(np.squeeze(g[\"label\"].values))\n            if y.ndim > 0 and len(y) >= 10: gs.append((x, y))\n        return gs\n    def _epoch(self, gs, train):\n        self.net.train(train)\n        if train: np.random.shuffle(gs)\n        total, cnt = 0.0, 0\n        if train: self.opt.zero_grad()\n        cm = torch.enable_grad() if train else torch.no_grad()\n        with cm:\n            for i, (x, y) in enumerate(gs):\n                xb = torch.FloatTensor(x).to(self.device)\n                yb = torch.FloatTensor(y).to(self.device)\n                pred = self.net(xb, mask_p=self.mask_p if train else 0.0)\n                raw = (pred - yb).pow(2).mean() - self.lam * self._ic(pred, yb)\n                if train:\n                    (raw / self.acc).backward()\n                    if (i + 1) % self.acc == 0 or i == len(gs) - 1:\n                        torch.nn.utils.clip_grad_value_(self.net.parameters(), 3.0)\n                        self.opt.step(); self.opt.zero_grad()\n                total += raw.item() * len(y); cnt += len(y)\n        return total / cnt\n    def fit(self, dataset):\n        df_t = dataset.prepare(\"train\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L)\n        df_v = dataset.prepare(\"valid\", col_set=[\"feature\", \"label\"], data_key=DataHandlerLP.DK_L)\n        tg, vg = self._group(df_t), self._group(df_v)\n        best_loss, best_ep = np.inf, 0\n        best_w = copy.deepcopy(self.net.state_dict())\n        for ep in range(self.n_epochs):\n            self._epoch(tg, True)\n            vl = self._epoch(vg, False)\n            if vl < best_loss:\n                best_loss, best_ep = vl, ep\n                best_w = copy.deepcopy(self.net.state_dict())\n            if ep - best_ep >= self.early_stop: break\n        self.net.load_state_dict(best_w)\n        self.fitted = True\n    def predict(self, dataset, segment=\"test\"):\n        if not self.fitted: raise ValueError(\"not fitted\")\n        x_df = dataset.prepare(segment, col_set=\"feature\", data_key=DataHandlerLP.DK_I)\n        x, bs = np.nan_to_num(x_df.values), 4096\n        self.net.eval()\n        parts = []\n        with torch.no_grad():\n            for s in range(0, len(x), bs):\n                xb = torch.FloatTensor(x[s:s + bs]).to(self.device)\n                parts.append(self.net(xb).reshape(-1).cpu().numpy())\n        return pd.Series(np.concatenate(parts), index=x_df.index)"
7}

1{
2  "n": 1
3}