mirror of
https://github.com/wassname/discovering_latent_knowledge.git
synced 2026-09-12 12:13:04 +08:00
mv from 023 to src, wip
This commit is contained in:
deep1
@@ -126,31 +126,7 @@
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def rows_item(row):\n",
|
||||
" \"\"\"\n",
|
||||
" transform a row by turning singe dim arrays into items\n",
|
||||
" \"\"\"\n",
|
||||
" for k,x in row.items():\n",
|
||||
" if isinstance(x, np.ndarray) and x.ndim==0:\n",
|
||||
" row[k]=x.item()\n",
|
||||
" return row\n",
|
||||
"\n",
|
||||
"def ds_info2df(ds):\n",
|
||||
" info = list(ds['info'])\n",
|
||||
" d = pd.DataFrame([rows_item(r) for r in info])\n",
|
||||
" return d\n",
|
||||
"\n",
|
||||
"def ds2df(ds):\n",
|
||||
" df = ds_info2df(ds)\n",
|
||||
" df_ans = ds.select_columns(['ans1', 'ans2', 'true', 'index', 'prob_y', 'prob_n', 'version']).with_format(\"numpy\").to_pandas()\n",
|
||||
" df = pd.concat([df, df_ans], axis=1)\n",
|
||||
" \n",
|
||||
" # derived\n",
|
||||
" df['dir_true'] = df['ans2'] - df['ans1']\n",
|
||||
" df['conf'] = (df['ans1']-df['ans2']).abs() \n",
|
||||
" df['llm_prob'] = (df['ans1']+df['ans2'])/2\n",
|
||||
" df['llm_ans'] = df['llm_prob']>0.5\n",
|
||||
" return df"
|
||||
"from src.datasets.load import ds2df"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -737,101 +713,8 @@
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def bool2switch(x):\n",
|
||||
" \"\"\"[0,1]->[-1,1]\"\"\"\n",
|
||||
" return x*2-1\n",
|
||||
"\n",
|
||||
"def switch2bool(x):\n",
|
||||
" \"\"\"[-1,1]->[0,1]\"\"\"\n",
|
||||
" return (x+1)/2\n",
|
||||
"\n",
|
||||
"assert switch2bool(-1)==0\n",
|
||||
"assert switch2bool(1)==1\n",
|
||||
"assert bool2switch(1)==1\n",
|
||||
"assert bool2switch(0)==-1\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def make_y(df):\n",
|
||||
" # label: is ans2 more true than ans1\n",
|
||||
" # so we ask does ans2 have greater probabiliy on \"positive\" than ans1\n",
|
||||
" # then, when the right answer is negative we swap the sign\n",
|
||||
" true_switch_sign = df.true_answer*2-1\n",
|
||||
" distance = (df.ans2-df.ans1) * true_switch_sign\n",
|
||||
" # y = bool2switch(distance>0)\n",
|
||||
" return distance"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"\n",
|
||||
"\n",
|
||||
"class imdbHSDataModule(pl.LightningDataModule):\n",
|
||||
"\n",
|
||||
" def __init__(self,\n",
|
||||
" ds,\n",
|
||||
" batch_size=32,\n",
|
||||
" ):\n",
|
||||
" super().__init__()\n",
|
||||
" self.save_hyperparameters(ignore=[\"ds\"])\n",
|
||||
" self.ds = ds.shuffle(seed=42)\n",
|
||||
"\n",
|
||||
" def setup(self, stage: str):\n",
|
||||
" h = self.hparams\n",
|
||||
" \n",
|
||||
" # extract data set into N-Dim tensors and 1-d dataframe\n",
|
||||
" self.ds_hs = (\n",
|
||||
" self.ds.select_columns(['hs1', 'hs2'])\n",
|
||||
" .with_format(\"numpy\")\n",
|
||||
" )\n",
|
||||
" self.df = ds2df(self.ds)\n",
|
||||
" \n",
|
||||
" y_cls = make_y(self.df)\n",
|
||||
" \n",
|
||||
" self.y = y_cls.values\n",
|
||||
" self.df['y'] = y_cls\n",
|
||||
" \n",
|
||||
" b = len(self.ds_hs)\n",
|
||||
" self.hs1 = self.ds_hs['hs1'].transpose(0, 2, 1)\n",
|
||||
" self.hs2 = self.ds_hs['hs2'].transpose(0, 2, 1)\n",
|
||||
" self.ans1 = self.df['ans1'].values\n",
|
||||
" self.ans2 = self.df['ans2'].values\n",
|
||||
"\n",
|
||||
" # let's create a simple 50/50 train split (the data is already randomized)\n",
|
||||
" n = len(self.y)\n",
|
||||
" \n",
|
||||
" self.val_split = vs = int(n * 0.5)\n",
|
||||
" self.test_split = ts = int(n * 0.75)\n",
|
||||
" hs1_train, hs2_train, y_train = self.hs1[:vs], self.hs2[:vs], self.y[:vs]\n",
|
||||
" hs1_val, hs2_val, y_val = self.hs1[vs:ts], self.hs2[vs:ts], self.y[vs:ts]\n",
|
||||
" hs1_test, hs2_test, y_test = self.hs1[ts:],self. hs2[ts:], self.y[ts:]\n",
|
||||
" \n",
|
||||
" \n",
|
||||
" to_ds = lambda x0, x1, y: TensorDataset(torch.from_numpy(x0).float(),\n",
|
||||
" torch.from_numpy(x1).float(),\n",
|
||||
" torch.from_numpy(y).float()\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" self.ds_train = to_ds(hs1_train, hs2_train, y_train)\n",
|
||||
"\n",
|
||||
" self.ds_val = to_ds(hs1_val, hs2_val, y_val)\n",
|
||||
"\n",
|
||||
" self.ds_test = to_ds(hs1_test, hs2_test, y_test)\n",
|
||||
"\n",
|
||||
" def train_dataloader(self):\n",
|
||||
" return DataLoader(self.ds_train,\n",
|
||||
" batch_size=self.hparams.batch_size,\n",
|
||||
" drop_last=True,\n",
|
||||
" shuffle=True)\n",
|
||||
"\n",
|
||||
" def val_dataloader(self):\n",
|
||||
" return DataLoader(self.ds_val, batch_size=self.hparams.batch_size, drop_last=True,)\n",
|
||||
"\n",
|
||||
" def test_dataloader(self):\n",
|
||||
" return DataLoader(self.ds_test, batch_size=self.hparams.batch_size, drop_last=True,)\n"
|
||||
"from src.helpers import switch2bool, bool2switch\n",
|
||||
"from src.datasets.dm import imdbHSDataModule"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -1001,125 +884,13 @@
|
||||
"# LightningModel"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 82,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"class MLPProbe(nn.Module):\n",
|
||||
" def __init__(self, c_in, depth=0, hs=16, dropout=0):\n",
|
||||
" super().__init__()\n",
|
||||
"\n",
|
||||
" layers = [\n",
|
||||
" nn.BatchNorm1d(c_in, affine=False), # this will normalise the inputs\n",
|
||||
" nn.Dropout1d(dropout),\n",
|
||||
" \n",
|
||||
" nn.Conv1d(c_in, hs*(depth+1), kernel_size=2),\n",
|
||||
" nn.ReLU(),\n",
|
||||
" nn.BatchNorm1d(hs*(depth+1)),\n",
|
||||
" ]\n",
|
||||
" for i in range(depth):\n",
|
||||
" layers += [\n",
|
||||
" nn.Conv1d(hs*(depth-i+1), hs*(depth-i), 2),\n",
|
||||
" nn.ReLU(),\n",
|
||||
" nn.BatchNorm1d(hs*(depth-i)),\n",
|
||||
" \n",
|
||||
" ]\n",
|
||||
" layers += [nn.AdaptiveAvgPool1d(1)]\n",
|
||||
" self.net = nn.Sequential(*layers)\n",
|
||||
" self.head = nn.Sequential(\n",
|
||||
" nn.Linear(hs, hs), nn.ReLU(),\n",
|
||||
" nn.Dropout(dropout), nn.Linear(hs, 1) \n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" def forward(self, x):\n",
|
||||
" h = self.net(x)\n",
|
||||
" # print(1, h.shape)\n",
|
||||
" h = h.squeeze(-1)\n",
|
||||
" # print(1, h.shape)\n",
|
||||
" return self.head(h)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 83,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from pytorch_optimizer import Ranger21\n",
|
||||
"import torchmetrics\n",
|
||||
"\n",
|
||||
"from torchmetrics import Metric, MetricCollection, Accuracy, AUROC\n",
|
||||
" \n",
|
||||
"class CSS(pl.LightningModule):\n",
|
||||
" def __init__(self, c_in, total_steps, depth=1, hs=16, lr=4e-3, weight_decay=1e-9, dropout=0):\n",
|
||||
" super().__init__()\n",
|
||||
" self.probe = MLPProbe(c_in, depth=depth, dropout=dropout, hs=hs)\n",
|
||||
" self.save_hyperparameters()\n",
|
||||
" \n",
|
||||
" self.loss_fn = nn.SmoothL1Loss()\n",
|
||||
" \n",
|
||||
" # metrics for each stage\n",
|
||||
" metrics_template = MetricCollection({\n",
|
||||
" 'acc': Accuracy(task=\"binary\"), \n",
|
||||
" 'auroc': AUROC(task=\"binary\")\n",
|
||||
" })\n",
|
||||
" self.metrics = torch.nn.ModuleDict({\n",
|
||||
" f'metrics_{stage}': metrics_template.clone(prefix=stage+'/') for stage in ['train', 'val', 'test']\n",
|
||||
" })\n",
|
||||
" \n",
|
||||
" def forward(self, x):\n",
|
||||
" return self.probe(x).squeeze(1)\n",
|
||||
" \n",
|
||||
" def _step(self, batch, batch_idx, stage='train'):\n",
|
||||
" x0, x1, y = batch\n",
|
||||
" ypred0 = self(x0)\n",
|
||||
" ypred1 = self(x1)\n",
|
||||
" \n",
|
||||
" if stage=='pred':\n",
|
||||
" return (ypred1-ypred0).float()\n",
|
||||
" \n",
|
||||
" loss = self.loss_fn(ypred1-ypred0, y)\n",
|
||||
" self.log(f\"{stage}/loss\", loss)\n",
|
||||
" \n",
|
||||
" m = self.metrics[f'metrics_{stage}']\n",
|
||||
" \n",
|
||||
" y_cls = switch2bool(ypred1-ypred0)\n",
|
||||
" m(y_cls, y>0.)\n",
|
||||
" self.log_dict(m, on_epoch=True, on_step=False)\n",
|
||||
" return loss\n",
|
||||
" \n",
|
||||
" def training_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
|
||||
" return self._step(batch, batch_idx)\n",
|
||||
" \n",
|
||||
" def validation_step(self, batch, batch_idx=0):\n",
|
||||
" return self._step(batch, batch_idx, stage='val')\n",
|
||||
" \n",
|
||||
" def predict_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
|
||||
" return self._step(batch, batch_idx, stage='pred').cpu().detach()\n",
|
||||
" \n",
|
||||
" def test_step(self, batch, batch_idx=0, dataloader_idx=0):\n",
|
||||
" return self._step(batch, batch_idx, stage='test')\n",
|
||||
" \n",
|
||||
" def configure_optimizers(self):\n",
|
||||
" \"\"\"use ranger21 from https://github.com/kozistr/pytorch_optimizer\"\"\"\n",
|
||||
" optimizer = Ranger21(\n",
|
||||
" self.parameters(),\n",
|
||||
" lr=self.hparams.lr,\n",
|
||||
" weight_decay=self.hparams.weight_decay, \n",
|
||||
" num_iterations=self.hparams.total_steps,\n",
|
||||
" )\n",
|
||||
" return optimizer\n",
|
||||
" \n",
|
||||
" "
|
||||
"from src.probs.conv import PLConvProbe"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -1245,7 +1016,7 @@
|
||||
"max_epochs = 42\n",
|
||||
"c_in = b[0].shape[1]\n",
|
||||
"print(b[0].shape)\n",
|
||||
"net = CSS(c_in=c_in, total_steps=max_epochs*len(dl_train), depth=6, hs=42*2, lr=3e-3, \n",
|
||||
"net = PLConvProbe(c_in=c_in, total_steps=max_epochs*len(dl_train), depth=6, hs=42*2, lr=3e-3, \n",
|
||||
" # weight_decay=1e-4, \n",
|
||||
" # dropout=0.1,\n",
|
||||
" )\n",
|
||||
@@ -2524,17 +2295,8 @@
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"# import pytorch_lightning as pl\n",
|
||||
"from lightning.pytorch.loggers.csv_logs import CSVLogger\n",
|
||||
"from pathlib import Path\n",
|
||||
"import pandas as pd\n",
|
||||
"from src.helpers.lightning import read_metrics_csv\n",
|
||||
"\n",
|
||||
"def read_metrics_csv(metrics_file_path):\n",
|
||||
" df_hist = pd.read_csv(metrics_file_path)\n",
|
||||
" df_hist[\"epoch\"] = df_hist[\"epoch\"].ffill()\n",
|
||||
" df_histe = df_hist.set_index(\"epoch\").groupby(\"epoch\").mean()\n",
|
||||
" return df_histe\n",
|
||||
" \n",
|
||||
"df_hist = read_metrics_csv(trainer.logger.experiment.metrics_file_path).ffill().bfill()\n",
|
||||
"df_hist"
|
||||
]
|
||||
|
||||
+1
-24
@@ -129,13 +129,6 @@
|
||||
"A uncensored and large one might be best for lying."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
@@ -362,23 +355,7 @@
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def ds_info2df(ds):\n",
|
||||
" d = pd.DataFrame(list(ds['info']))\n",
|
||||
" # for c in ['desired_answer', 'lie', 'true_answer']:\n",
|
||||
" # d[c] = d[c].map(lambda x:x.item())\n",
|
||||
" return d\n",
|
||||
"\n",
|
||||
"def ds2df(ds):\n",
|
||||
" df = ds_info2df(ds)\n",
|
||||
" df_ans = ds.select_columns(['ans1', 'ans2', 'true', 'index', 'version']).with_format(\"numpy\").to_pandas()\n",
|
||||
" df = pd.concat([df, df_ans], axis=1)\n",
|
||||
" \n",
|
||||
" # derived\n",
|
||||
" df['dir_true'] = df['ans2'] - df['ans1']\n",
|
||||
" df['conf'] = (df['ans1']-df['ans2']).abs() \n",
|
||||
" df['llm_prob'] = (df['ans1']+df['ans2'])/2\n",
|
||||
" df['llm_ans'] = df['llm_prob']>0.5\n",
|
||||
" return df\n"
|
||||
"from src.datasets.load import ds2df"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import lightning as pl
|
||||
import pandas as pd
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
from src.datasets.load import ds2df
|
||||
|
||||
def make_y(df):
|
||||
# label: is ans2 more true than ans1
|
||||
# so we ask does ans2 have greater probabiliy on "positive" than ans1
|
||||
# then, when the right answer is negative we swap the sign
|
||||
true_switch_sign = df.true_answer*2-1
|
||||
distance = (df.ans2-df.ans1) * true_switch_sign
|
||||
# y = bool2switch(distance>0)
|
||||
return distance
|
||||
|
||||
class imdbHSDataModule(pl.LightningDataModule):
|
||||
|
||||
def __init__(self,
|
||||
ds: Dataset,
|
||||
batch_size: int=32,
|
||||
):
|
||||
super().__init__()
|
||||
self.save_hyperparameters(ignore=["ds"])
|
||||
self.ds = ds.shuffle(seed=42)
|
||||
|
||||
def setup(self, stage: str):
|
||||
h = self.hparams
|
||||
|
||||
# extract data set into N-Dim tensors and 1-d dataframe
|
||||
self.ds_hs = (
|
||||
self.ds.select_columns(['hs1', 'hs2'])
|
||||
.with_format("numpy")
|
||||
)
|
||||
self.df = ds2df(self.ds)
|
||||
|
||||
y_cls = make_y(self.df)
|
||||
|
||||
self.y = y_cls.values
|
||||
self.df['y'] = y_cls
|
||||
|
||||
b = len(self.ds_hs)
|
||||
self.hs1 = self.ds_hs['hs1'].transpose(0, 2, 1)
|
||||
self.hs2 = self.ds_hs['hs2'].transpose(0, 2, 1)
|
||||
self.ans1 = self.df['ans1'].values
|
||||
self.ans2 = self.df['ans2'].values
|
||||
|
||||
# let's create a simple 50/50 train split (the data is already randomized)
|
||||
n = len(self.y)
|
||||
|
||||
self.val_split = vs = int(n * 0.5)
|
||||
self.test_split = ts = int(n * 0.75)
|
||||
hs1_train, hs2_train, y_train = self.hs1[:vs], self.hs2[:vs], self.y[:vs]
|
||||
hs1_val, hs2_val, y_val = self.hs1[vs:ts], self.hs2[vs:ts], self.y[vs:ts]
|
||||
hs1_test, hs2_test, y_test = self.hs1[ts:],self. hs2[ts:], self.y[ts:]
|
||||
|
||||
|
||||
to_ds = lambda x0, x1, y: TensorDataset(torch.from_numpy(x0).float(),
|
||||
torch.from_numpy(x1).float(),
|
||||
torch.from_numpy(y).float()
|
||||
)
|
||||
|
||||
self.ds_train = to_ds(hs1_train, hs2_train, y_train)
|
||||
|
||||
self.ds_val = to_ds(hs1_val, hs2_val, y_val)
|
||||
|
||||
self.ds_test = to_ds(hs1_test, hs2_test, y_test)
|
||||
|
||||
def train_dataloader(self):
|
||||
return DataLoader(self.ds_train,
|
||||
batch_size=self.hparams.batch_size,
|
||||
drop_last=True,
|
||||
shuffle=True)
|
||||
|
||||
def val_dataloader(self):
|
||||
return DataLoader(self.ds_val, batch_size=self.hparams.batch_size, drop_last=True,)
|
||||
|
||||
def test_dataloader(self):
|
||||
return DataLoader(self.ds_test, batch_size=self.hparams.batch_size, drop_last=True,)
|
||||
@@ -0,0 +1,25 @@
|
||||
def rows_item(row):
|
||||
"""
|
||||
transform a row by turning singe dim arrays into items
|
||||
"""
|
||||
for k,x in row.items():
|
||||
if isinstance(x, np.ndarray) and x.ndim==0:
|
||||
row[k]=x.item()
|
||||
return row
|
||||
|
||||
def ds_info2df(ds):
|
||||
info = list(ds['info'])
|
||||
d = pd.DataFrame([rows_item(r) for r in info])
|
||||
return d
|
||||
|
||||
def ds2df(ds):
|
||||
df = ds_info2df(ds)
|
||||
df_ans = ds.select_columns(['ans1', 'ans2', 'true', 'index', 'prob_y', 'prob_n', 'version']).with_format("numpy").to_pandas()
|
||||
df = pd.concat([df, df_ans], axis=1)
|
||||
|
||||
# derived
|
||||
df['dir_true'] = df['ans2'] - df['ans1']
|
||||
df['conf'] = (df['ans1']-df['ans2']).abs()
|
||||
df['llm_prob'] = (df['ans1']+df['ans2'])/2
|
||||
df['llm_ans'] = df['llm_prob']>0.5
|
||||
return df
|
||||
@@ -0,0 +1,7 @@
|
||||
def bool2switch(x):
|
||||
"""[0,1]->[-1,1]"""
|
||||
return x*2-1
|
||||
|
||||
def switch2bool(x):
|
||||
"""[-1,1]->[0,1]"""
|
||||
return (x+1)/2
|
||||
@@ -0,0 +1,10 @@
|
||||
from lightning.pytorch.loggers.csv_logs import CSVLogger
|
||||
from pathlib import Path
|
||||
import pandas as pd
|
||||
|
||||
def read_metrics_csv(metrics_file_path):
|
||||
df_hist = pd.read_csv(metrics_file_path)
|
||||
df_hist["epoch"] = df_hist["epoch"].ffill()
|
||||
df_histe = df_hist.set_index("epoch").groupby("epoch").mean()
|
||||
return df_histe
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .pl_ranking import PLRanking
|
||||
|
||||
class ConvProbe(nn.Module):
|
||||
def __init__(self, c_in, depth=0, hs=16, dropout=0):
|
||||
super().__init__()
|
||||
|
||||
layers = [
|
||||
nn.BatchNorm1d(c_in, affine=False), # this will normalise the inputs
|
||||
nn.Dropout1d(dropout),
|
||||
|
||||
nn.Conv1d(c_in, hs*(depth+1), kernel_size=2),
|
||||
nn.ReLU(),
|
||||
nn.BatchNorm1d(hs*(depth+1)),
|
||||
]
|
||||
for i in range(depth):
|
||||
layers += [
|
||||
nn.Conv1d(hs*(depth-i+1), hs*(depth-i), 2),
|
||||
nn.ReLU(),
|
||||
nn.BatchNorm1d(hs*(depth-i)),
|
||||
|
||||
]
|
||||
layers += [nn.AdaptiveAvgPool1d(1)]
|
||||
self.net = nn.Sequential(*layers)
|
||||
self.head = nn.Sequential(
|
||||
nn.Linear(hs, hs), nn.ReLU(),
|
||||
nn.Dropout(dropout), nn.Linear(hs, 1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
h = self.net(x)
|
||||
# print(1, h.shape)
|
||||
h = h.squeeze(-1)
|
||||
# print(1, h.shape)
|
||||
return self.head(h)
|
||||
|
||||
|
||||
class PLConvProbe(PLRanking):
|
||||
def __init__(self, *args, **kwargs)
|
||||
super().__init__(*args, **kwargs)
|
||||
self.probe = MLPProbe(c_in, depth=depth, dropout=dropout, hs=hs)
|
||||
@@ -0,0 +1,72 @@
|
||||
from pytorch_optimizer import Ranger21
|
||||
import torchmetrics
|
||||
from src.helpers import switch2bool, bool2switch
|
||||
|
||||
from torchmetrics import Metric, MetricCollection, Accuracy, AUROC
|
||||
|
||||
class PLRanking(pl.LightningModule):
|
||||
"""
|
||||
Base pytorch lightning module, subclass to add model
|
||||
|
||||
This uses SmoothL1Loss to tackle a ranking objective and does better in terms of performance and overfitting compared to MarginRanking loss, as well as setting it up to classify the direction, or estimate the distance between the pair direction with MSE.
|
||||
"""
|
||||
def __init__(self, c_in, total_steps, depth=1, hs=16, lr=4e-3, weight_decay=1e-9, dropout=0):
|
||||
super().__init__()
|
||||
# self.probe = MLPProbe(c_in, depth=depth, dropout=dropout, hs=hs)
|
||||
self.save_hyperparameters()
|
||||
|
||||
self.loss_fn = nn.SmoothL1Loss()
|
||||
|
||||
# metrics for each stage
|
||||
metrics_template = MetricCollection({
|
||||
'acc': Accuracy(task="binary"),
|
||||
'auroc': AUROC(task="binary")
|
||||
})
|
||||
self.metrics = torch.nn.ModuleDict({
|
||||
f'metrics_{stage}': metrics_template.clone(prefix=stage+'/') for stage in ['train', 'val', 'test']
|
||||
})
|
||||
|
||||
def forward(self, x):
|
||||
return self.probe(x).squeeze(1)
|
||||
|
||||
def _step(self, batch, batch_idx, stage='train'):
|
||||
x0, x1, y = batch
|
||||
ypred0 = self(x0)
|
||||
ypred1 = self(x1)
|
||||
|
||||
if stage=='pred':
|
||||
return (ypred1-ypred0).float()
|
||||
|
||||
loss = self.loss_fn(ypred1-ypred0, y)
|
||||
self.log(f"{stage}/loss", loss)
|
||||
|
||||
m = self.metrics[f'metrics_{stage}']
|
||||
|
||||
y_cls = switch2bool(ypred1-ypred0)
|
||||
m(y_cls, y>0.)
|
||||
self.log_dict(m, on_epoch=True, on_step=False)
|
||||
return loss
|
||||
|
||||
def training_step(self, batch, batch_idx=0, dataloader_idx=0):
|
||||
return self._step(batch, batch_idx)
|
||||
|
||||
def validation_step(self, batch, batch_idx=0):
|
||||
return self._step(batch, batch_idx, stage='val')
|
||||
|
||||
def predict_step(self, batch, batch_idx=0, dataloader_idx=0):
|
||||
return self._step(batch, batch_idx, stage='pred').cpu().detach()
|
||||
|
||||
def test_step(self, batch, batch_idx=0, dataloader_idx=0):
|
||||
return self._step(batch, batch_idx, stage='test')
|
||||
|
||||
def configure_optimizers(self):
|
||||
"""use ranger21 from https://github.com/kozistr/pytorch_optimizer"""
|
||||
optimizer = Ranger21(
|
||||
self.parameters(),
|
||||
lr=self.hparams.lr,
|
||||
weight_decay=self.hparams.weight_decay,
|
||||
num_iterations=self.hparams.total_steps,
|
||||
)
|
||||
return optimizer
|
||||
|
||||
|
||||
Reference in New Issue
Block a user