hyena predict longer sequences
This commit is contained in:
martin committed
1 parent
bdef020bff
commit
a71030547c
4 files changed
+1980
-72
No files matched your search
@@ -12,7 +12,7 @@
|
|||||||
"text": [
|
"text": [
|
||||||
"torch.Size([1, 1024, 128]) torch.Size([1, 1024, 128])\n",
|
"torch.Size([1, 1024, 128]) torch.Size([1, 1024, 128])\n",
|
||||||
"Causality check: gradients should not flow \"from future to past\"\n",
|
"Causality check: gradients should not flow \"from future to past\"\n",
|
||||||
"tensor(2.1330e-09) tensor(0.2463)\n"
|
"tensor(-4.1268e-09) tensor(0.0844)\n"
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
@@ -291,7 +291,7 @@
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
"cell_type": "code",
|
"cell_type": "code",
|
||||||
"execution_count": 4,
|
"execution_count": 2,
|
||||||
"id": "032ef08a-8cc6-491a-9eb8-4a6b3f2d165e",
|
"id": "032ef08a-8cc6-491a-9eb8-4a6b3f2d165e",
|
||||||
"metadata": {},
|
"metadata": {},
|
||||||
"outputs": [
|
"outputs": [
|
||||||
@@ -343,7 +343,9 @@
|
|||||||
" # Increase the channel dimension from 1 to d_model\n",
|
" # Increase the channel dimension from 1 to d_model\n",
|
||||||
" u = self.fc_before(u) \n",
|
" u = self.fc_before(u) \n",
|
||||||
" # Pass through the operator\n",
|
" # Pass through the operator\n",
|
||||||
|
" #[B,1024,128] --> [B,1024,128]\n",
|
||||||
" u = self.operator(u)\n",
|
" u = self.operator(u)\n",
|
||||||
|
" \n",
|
||||||
" last_state = u[:,-1,:]\n",
|
" last_state = u[:,-1,:]\n",
|
||||||
" # Decrease the channel dimension back to 1\n",
|
" # Decrease the channel dimension back to 1\n",
|
||||||
" y = self.fc_after(last_state)\n",
|
" y = self.fc_after(last_state)\n",
|
||||||
@@ -371,7 +373,7 @@
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
"cell_type": "code",
|
"cell_type": "code",
|
||||||
"execution_count": 10,
|
"execution_count": 3,
|
||||||
"id": "80cde67b-992f-4cb0-8824-4a6b7e4984ca",
|
"id": "80cde67b-992f-4cb0-8824-4a6b7e4984ca",
|
||||||
"metadata": {},
|
"metadata": {},
|
||||||
"outputs": [
|
"outputs": [
|
||||||
@@ -379,16 +381,16 @@
|
|||||||
"name": "stdout",
|
"name": "stdout",
|
||||||
"output_type": "stream",
|
"output_type": "stream",
|
||||||
"text": [
|
"text": [
|
||||||
"Train Epoch: 1 [0/640 (0%)]\tLoss: 0.736030\n",
|
"Train Epoch: 1 [0/640 (0%)]\tLoss: 0.446847\n",
|
||||||
"Train Epoch: 2 [0/640 (0%)]\tLoss: 0.013385\n",
|
"Train Epoch: 2 [0/640 (0%)]\tLoss: 0.077979\n",
|
||||||
"Train Epoch: 3 [0/640 (0%)]\tLoss: 0.019001\n",
|
"Train Epoch: 3 [0/640 (0%)]\tLoss: 0.021656\n",
|
||||||
"Train Epoch: 4 [0/640 (0%)]\tLoss: 0.010262\n",
|
"Train Epoch: 4 [0/640 (0%)]\tLoss: 0.007355\n",
|
||||||
"Train Epoch: 5 [0/640 (0%)]\tLoss: 0.005347\n",
|
"Train Epoch: 5 [0/640 (0%)]\tLoss: 0.004926\n",
|
||||||
"Train Epoch: 6 [0/640 (0%)]\tLoss: 0.006345\n",
|
"Train Epoch: 6 [0/640 (0%)]\tLoss: 0.006014\n",
|
||||||
"Train Epoch: 7 [0/640 (0%)]\tLoss: 0.004454\n",
|
"Train Epoch: 7 [0/640 (0%)]\tLoss: 0.003400\n",
|
||||||
"Train Epoch: 8 [0/640 (0%)]\tLoss: 0.003857\n",
|
"Train Epoch: 8 [0/640 (0%)]\tLoss: 0.003720\n",
|
||||||
"Train Epoch: 9 [0/640 (0%)]\tLoss: 0.003062\n",
|
"Train Epoch: 9 [0/640 (0%)]\tLoss: 0.004267\n",
|
||||||
"Train Epoch: 10 [0/640 (0%)]\tLoss: 0.002607\n"
|
"Train Epoch: 10 [0/640 (0%)]\tLoss: 0.004081\n"
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
@@ -495,49 +497,155 @@
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
"cell_type": "code",
|
"cell_type": "code",
|
||||||
"execution_count": 17,
|
"execution_count": 4,
|
||||||
"id": "1b763e03-baab-4b02-bae0-5747461bca7f",
|
"id": "90330622-8b44-4b45-8158-6840538f768c",
|
||||||
"metadata": {},
|
"metadata": {},
|
||||||
"outputs": [],
|
"outputs": [],
|
||||||
"source": [
|
"source": [
|
||||||
"def correlate(model, device, data_loader):\n",
|
"from sklearn.linear_model import LinearRegression\n",
|
||||||
" model.eval()\n",
|
"from sklearn.metrics import r2_score\n",
|
||||||
" correlations = {key: [] for key in [\"frequency\", \"phase\", \"amplitude\", \"noise_sd\"]}\n",
|
|
||||||
" with torch.no_grad():\n",
|
|
||||||
" for data, target, params in data_loader:\n",
|
|
||||||
"\n",
|
"\n",
|
||||||
" data, target = data.to(device), target.to(device)\n",
|
"def fit_and_evaluate_linear_regression(outputs_and_params):\n",
|
||||||
" output, last_state = model(data)\n",
|
" # Split the data into inputs (last_states) and targets (params)\n",
|
||||||
" last_state_np = last_state.cpu().numpy()\n",
|
" inputs = np.concatenate([x[0] for x in outputs_and_params])\n",
|
||||||
" params_np = params.cpu().numpy()\n",
|
" targets = np.concatenate([x[1] for x in outputs_and_params])\n",
|
||||||
" #import pdb;pdb.set_trace()\n",
|
|
||||||
" \n",
|
" \n",
|
||||||
" # Compute correlations between last_state and parameters\n",
|
" r2_scores = []\n",
|
||||||
" for i, key in enumerate(correlations.keys()):\n",
|
" param_names = [\"frequency\", \"phase\", \"amplitude\", \"noise_sd\"]\n",
|
||||||
" correlation = np.corrcoef(last_state_np.squeeze(), params_np[:,i])[0,1]\n",
|
|
||||||
" correlations[key].append(correlation)\n",
|
|
||||||
" \n",
|
" \n",
|
||||||
" # Average correlations over all batches\n",
|
" # Fit the linear regression model for each parameter and calculate the R^2 score\n",
|
||||||
" avg_correlations = {key: np.mean(value) for key, value in correlations.items()}\n",
|
" for i in range(targets.shape[1]):\n",
|
||||||
" return avg_correlations"
|
" model = LinearRegression().fit(inputs, targets[:, i])\n",
|
||||||
|
" pred = model.predict(inputs)\n",
|
||||||
|
" score = r2_score(targets[:, i], pred)\n",
|
||||||
|
" r2_scores.append(score)\n",
|
||||||
|
" print(f\"R^2 score for {param_names[i]}: {score:.2f}\")\n",
|
||||||
|
" \n",
|
||||||
|
" return r2_scores"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"cell_type": "code",
|
"cell_type": "code",
|
||||||
"execution_count": 18,
|
"execution_count": 5,
|
||||||
"id": "f4c78c51-a538-4d24-ab7b-fb6035c78df8",
|
"id": "5eb62a22-cad8-43c4-b757-f36b6a01e9be",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def generate_outputs(model, device, data_loader):\n",
|
||||||
|
" model.eval()\n",
|
||||||
|
" outputs_and_params = []\n",
|
||||||
|
" with torch.no_grad():\n",
|
||||||
|
" for data, target, params in data_loader:\n",
|
||||||
|
" data, target = data.to(device), target.to(device)\n",
|
||||||
|
" output, last_state = model(data)\n",
|
||||||
|
" outputs_and_params.append((last_state.cpu().numpy(), params.cpu().numpy()))\n",
|
||||||
|
" return outputs_and_params\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 6,
|
||||||
|
"id": "a95ee542-1c39-4f04-9184-e26c6983a018",
|
||||||
"metadata": {},
|
"metadata": {},
|
||||||
"outputs": [
|
"outputs": [
|
||||||
{
|
{
|
||||||
"name": "stdout",
|
"name": "stdout",
|
||||||
"output_type": "stream",
|
"output_type": "stream",
|
||||||
"text": [
|
"text": [
|
||||||
"> \u001b[0;32m/tmp/ipykernel_9454/3342195123.py\u001b[0m(14)\u001b[0;36mcorrelate\u001b[0;34m()\u001b[0m\n",
|
"R^2 score for frequency: 0.77\n",
|
||||||
"\u001b[0;32m 12 \u001b[0;31m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
"R^2 score for phase: 0.66\n",
|
||||||
"\u001b[0m\u001b[0;32m 13 \u001b[0;31m \u001b[0;31m# Compute correlations between last_state and parameters\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
"R^2 score for amplitude: 0.99\n",
|
||||||
"\u001b[0m\u001b[0;32m---> 14 \u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkey\u001b[0m \u001b[0;32min\u001b[0m \u001b[0menumerate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcorrelations\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mkeys\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
"R^2 score for noise_sd: 0.97\n"
|
||||||
"\u001b[0m\u001b[0;32m 15 \u001b[0;31m \u001b[0mcorrelation\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcorrcoef\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlast_state_np\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams_np\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
]
|
||||||
"\u001b[0m\u001b[0;32m 16 \u001b[0;31m \u001b[0mcorrelations\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcorrelation\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"outputs_and_params = generate_outputs(model, device, train_loader)\n",
|
||||||
|
"\n",
|
||||||
|
"# Fit the linear regression model and print the R^2 score for each parameter\n",
|
||||||
|
"r2_scores = fit_and_evaluate_linear_regression(outputs_and_params)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "de65d0d2-b0c6-4ac5-a87f-70fa7f90480b",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"# Autoregressive"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 22,
|
||||||
|
"id": "8ee139a6-aee5-4309-8685-bf0c28893279",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def predict_autoregressive(model, initial_data, n_steps):\n",
|
||||||
|
" model.eval()\n",
|
||||||
|
" predictions = []\n",
|
||||||
|
" current_input = initial_data\n",
|
||||||
|
" with torch.no_grad():\n",
|
||||||
|
" for _ in range(n_steps):\n",
|
||||||
|
" # Get the prediction for the next step and save it\n",
|
||||||
|
" import pdb;pdb.set_trace()\n",
|
||||||
|
" next_output, _ = model(current_input)\n",
|
||||||
|
" predictions.append(next_output)\n",
|
||||||
|
"\n",
|
||||||
|
" # Prepare the input for the next step\n",
|
||||||
|
" next_input = torch.cat((current_input[:, 1:, :], next_output[:, -1:, :]), dim=1)\n",
|
||||||
|
"\n",
|
||||||
|
" current_input = next_input\n",
|
||||||
|
"\n",
|
||||||
|
" return torch.cat(predictions, dim=1)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "272040d6-4dfb-438b-a432-744e950effd1",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 23,
|
||||||
|
"id": "9f49cbf8-08e8-428a-af80-bcb2515a6327",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"torch.Size([1, 1024, 1])"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 23,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"initial_data = dataset[0]\n",
|
||||||
|
"initial_data[0][None,...].shape"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 24,
|
||||||
|
"id": "c75464e3-e44d-4325-8804-29d8081e3a45",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"> \u001b[0;32m/tmp/ipykernel_9626/3477884695.py\u001b[0m(9)\u001b[0;36mpredict_autoregressive\u001b[0;34m()\u001b[0m\n",
|
||||||
|
"\u001b[0;32m 7 \u001b[0;31m \u001b[0;31m# Get the prediction for the next step and save it\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
||||||
|
"\u001b[0m\u001b[0;32m 8 \u001b[0;31m \u001b[0;32mimport\u001b[0m \u001b[0mpdb\u001b[0m\u001b[0;34m;\u001b[0m\u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
||||||
|
"\u001b[0m\u001b[0;32m----> 9 \u001b[0;31m \u001b[0mnext_output\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcurrent_input\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
||||||
|
"\u001b[0m\u001b[0;32m 10 \u001b[0;31m \u001b[0mpredictions\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnext_output\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
||||||
|
"\u001b[0m\u001b[0;32m 11 \u001b[0;31m\u001b[0;34m\u001b[0m\u001b[0m\n",
|
||||||
"\u001b[0m\n"
|
"\u001b[0m\n"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
@@ -545,53 +653,34 @@
|
|||||||
"name": "stdin",
|
"name": "stdin",
|
||||||
"output_type": "stream",
|
"output_type": "stream",
|
||||||
"text": [
|
"text": [
|
||||||
"ipdb> last_state_np.shape\n"
|
"ipdb> current_input\n"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"name": "stdout",
|
"name": "stdout",
|
||||||
"output_type": "stream",
|
"output_type": "stream",
|
||||||
"text": [
|
"text": [
|
||||||
"(64, 128)\n"
|
"tensor([[-1.2113],\n",
|
||||||
|
" [-1.2368],\n",
|
||||||
|
" [-1.1851],\n",
|
||||||
|
" ...,\n",
|
||||||
|
" [-1.1084],\n",
|
||||||
|
" [-0.9648],\n",
|
||||||
|
" [-1.0586]])\n"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"name": "stdin",
|
"name": "stdin",
|
||||||
"output_type": "stream",
|
"output_type": "stream",
|
||||||
"text": [
|
"text": [
|
||||||
"ipdb> last_state_np[0]\n"
|
"ipdb> current_input.shape\n"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"name": "stdout",
|
"name": "stdout",
|
||||||
"output_type": "stream",
|
"output_type": "stream",
|
||||||
"text": [
|
"text": [
|
||||||
"array([ 0.13440676, -0.03606963, -0.4140934 , 0.5431792 , 0.36556095,\n",
|
"torch.Size([1024, 1])\n"
|
||||||
" 0.2256596 , -0.4616309 , -0.05567896, 0.17625177, -0.23529659,\n",
|
|
||||||
" -0.5208519 , 0.29691923, 0.15615058, 0.31342992, -0.5054718 ,\n",
|
|
||||||
" -0.33130994, -0.03956199, 0.31403548, -0.15925817, -0.22006416,\n",
|
|
||||||
" 0.00838468, -0.30691615, -0.1828884 , -0.52498204, 0.07198659,\n",
|
|
||||||
" 0.38572663, -0.27560705, 0.12110637, -0.17199083, 0.3913066 ,\n",
|
|
||||||
" -0.03978934, -0.21167544, -0.43025637, 0.20562531, 0.3000516 ,\n",
|
|
||||||
" -0.6784174 , -0.04233613, 0.4706083 , 0.20292807, 0.49932548,\n",
|
|
||||||
" 0.00203749, 0.2665777 , -0.16989222, 0.40648764, 0.22203793,\n",
|
|
||||||
" -0.44289762, 0.20751204, -0.38801843, -0.001487 , -0.49365598,\n",
|
|
||||||
" 0.05991718, -0.10120638, 0.36523518, -0.15450253, 0.11142011,\n",
|
|
||||||
" -0.20295474, 0.12229299, 0.09449576, -0.3422598 , 0.18969077,\n",
|
|
||||||
" 0.517254 , 0.08046471, 0.02134303, -0.35802346, -0.26192123,\n",
|
|
||||||
" 0.26145002, 0.11439252, 0.03314593, -0.15331428, 0.42282102,\n",
|
|
||||||
" 0.6026961 , -0.04233361, -0.5652172 , 0.33544067, 0.05744396,\n",
|
|
||||||
" 0.43544483, 0.2176097 , 0.22265801, -0.03894311, -0.01405966,\n",
|
|
||||||
" 0.23479447, 0.32931918, 0.21597862, 0.40402904, 0.20630498,\n",
|
|
||||||
" 0.09036086, -0.16922598, -0.1774486 , -0.14753146, -0.22214624,\n",
|
|
||||||
" -0.19101782, -0.09274255, 0.10928088, -0.01354241, -0.3864469 ,\n",
|
|
||||||
" 0.46331462, -0.38134843, -0.07766411, 0.750954 , -0.06306303,\n",
|
|
||||||
" -0.33691666, -0.1798551 , 0.19826202, -0.13544285, 0.01956506,\n",
|
|
||||||
" 0.6431204 , -0.11272874, 0.1345196 , 0.23029736, 0.28865197,\n",
|
|
||||||
" 0.70087713, 0.3592593 , 0.30329305, -0.26943353, -0.11942452,\n",
|
|
||||||
" -0.21187985, 0.19452253, 0.05659255, -0.00958484, -0.33417243,\n",
|
|
||||||
" -0.14836329, 0.28580692, 0.20885246, 0.18010336, 0.56253076,\n",
|
|
||||||
" -0.25303417, 0.0189368 , 0.2504725 ], dtype=float32)\n"
|
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
@@ -603,13 +692,13 @@
|
|||||||
}
|
}
|
||||||
],
|
],
|
||||||
"source": [
|
"source": [
|
||||||
"correlate(model,\"cpu\",train_loader)"
|
"predict_autoregressive(model,initial_data[0][None,...],100)"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
"cell_type": "code",
|
"cell_type": "code",
|
||||||
"execution_count": null,
|
"execution_count": null,
|
||||||
"id": "e1558ffb-4699-4a8c-b05d-c3ac31a3829f",
|
"id": "90a37c56-59f3-49bc-bfbe-d9e126c42ed1",
|
||||||
"metadata": {},
|
"metadata": {},
|
||||||
"outputs": [],
|
"outputs": [],
|
||||||
"source": []
|
"source": []
|
||||||
|
|||||||
@@ -0,0 +1,850 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 32,
|
||||||
|
"id": "6c6e33cb-72f9-42fa-936a-33b5fe338d15",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"torch.Size([1, 1024, 128]) torch.Size([1, 1024, 128])\n",
|
||||||
|
"Causality check: gradients should not flow \"from future to past\"\n",
|
||||||
|
"tensor(-6.7998e-10) tensor(0.1017)\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"# %load standalone_hyena.py\n",
|
||||||
|
"\"\"\"\n",
|
||||||
|
"Simplified standalone version of Hyena: https://arxiv.org/abs/2302.10866, designed for quick experimentation.\n",
|
||||||
|
"A complete version is available under `src.models.sequence.hyena`.\n",
|
||||||
|
"\"\"\"\n",
|
||||||
|
"\n",
|
||||||
|
"import math\n",
|
||||||
|
"import torch\n",
|
||||||
|
"import torch.nn as nn\n",
|
||||||
|
"import torch.nn.functional as F\n",
|
||||||
|
"from einops import rearrange\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def fftconv(u, k, D):\n",
|
||||||
|
" seqlen = u.shape[-1]\n",
|
||||||
|
" fft_size = 2 * seqlen\n",
|
||||||
|
" \n",
|
||||||
|
" k_f = torch.fft.rfft(k, n=fft_size) / fft_size\n",
|
||||||
|
" u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)\n",
|
||||||
|
" \n",
|
||||||
|
" if len(u.shape) > 3: k_f = k_f.unsqueeze(1)\n",
|
||||||
|
" y = torch.fft.irfft(u_f * k_f, n=fft_size, norm='forward')[..., :seqlen]\n",
|
||||||
|
"\n",
|
||||||
|
" out = y + u * D.unsqueeze(-1)\n",
|
||||||
|
" return out.to(dtype=u.dtype)\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"@torch.jit.script \n",
|
||||||
|
"def mul_sum(q, y):\n",
|
||||||
|
" return (q * y).sum(dim=1)\n",
|
||||||
|
"\n",
|
||||||
|
"class OptimModule(nn.Module):\n",
|
||||||
|
" \"\"\" Interface for Module that allows registering buffers/parameters with configurable optimizer hyperparameters \"\"\"\n",
|
||||||
|
"\n",
|
||||||
|
" def register(self, name, tensor, lr=None, wd=0.0):\n",
|
||||||
|
" \"\"\"Register a tensor with a configurable learning rate and 0 weight decay\"\"\"\n",
|
||||||
|
"\n",
|
||||||
|
" if lr == 0.0:\n",
|
||||||
|
" self.register_buffer(name, tensor)\n",
|
||||||
|
" else:\n",
|
||||||
|
" self.register_parameter(name, nn.Parameter(tensor))\n",
|
||||||
|
"\n",
|
||||||
|
" optim = {}\n",
|
||||||
|
" if lr is not None: optim[\"lr\"] = lr\n",
|
||||||
|
" if wd is not None: optim[\"weight_decay\"] = wd\n",
|
||||||
|
" setattr(getattr(self, name), \"_optim\", optim)\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
"class Sin(nn.Module):\n",
|
||||||
|
" def __init__(self, dim, w=10, train_freq=True):\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" self.freq = nn.Parameter(w * torch.ones(1, dim)) if train_freq else w * torch.ones(1, dim)\n",
|
||||||
|
"\n",
|
||||||
|
" def forward(self, x):\n",
|
||||||
|
" return torch.sin(self.freq * x)\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
"class PositionalEmbedding(OptimModule):\n",
|
||||||
|
" def __init__(self, emb_dim: int, seq_len: int, lr_pos_emb: float=1e-5, **kwargs): \n",
|
||||||
|
" \"\"\"Complex exponential positional embeddings for Hyena filters.\"\"\" \n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" \n",
|
||||||
|
" self.seq_len = seq_len\n",
|
||||||
|
" # The time embedding fed to the filteres is normalized so that t_f = 1\n",
|
||||||
|
" t = torch.linspace(0, 1, self.seq_len)[None, :, None] # 1, L, 1\n",
|
||||||
|
" \n",
|
||||||
|
" if emb_dim > 1:\n",
|
||||||
|
" bands = (emb_dim - 1) // 2 \n",
|
||||||
|
" # To compute the right embeddings we use the \"proper\" linspace \n",
|
||||||
|
" t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]\n",
|
||||||
|
" w = 2 * math.pi * t_rescaled / seq_len # 1, L, 1 \n",
|
||||||
|
" \n",
|
||||||
|
" f = torch.linspace(1e-4, bands - 1, bands)[None, None] \n",
|
||||||
|
" z = torch.exp(-1j * f * w)\n",
|
||||||
|
" z = torch.cat([t, z.real, z.imag], dim=-1)\n",
|
||||||
|
" self.register(\"z\", z, lr=lr_pos_emb) \n",
|
||||||
|
" self.register(\"t\", t, lr=0.0)\n",
|
||||||
|
" \n",
|
||||||
|
" def forward(self, L):\n",
|
||||||
|
" return self.z[:, :L], self.t[:, :L]\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
"class ExponentialModulation(OptimModule):\n",
|
||||||
|
" def __init__(\n",
|
||||||
|
" self,\n",
|
||||||
|
" d_model,\n",
|
||||||
|
" fast_decay_pct=0.3,\n",
|
||||||
|
" slow_decay_pct=1.5,\n",
|
||||||
|
" target=1e-2,\n",
|
||||||
|
" modulation_lr=0.0,\n",
|
||||||
|
" modulate: bool=True,\n",
|
||||||
|
" shift: float = 0.0,\n",
|
||||||
|
" **kwargs\n",
|
||||||
|
" ):\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" self.modulate = modulate\n",
|
||||||
|
" self.shift = shift\n",
|
||||||
|
" max_decay = math.log(target) / fast_decay_pct\n",
|
||||||
|
" min_decay = math.log(target) / slow_decay_pct\n",
|
||||||
|
" deltas = torch.linspace(min_decay, max_decay, d_model)[None, None]\n",
|
||||||
|
" self.register(\"deltas\", deltas, lr=modulation_lr)\n",
|
||||||
|
" \n",
|
||||||
|
" def forward(self, t, x):\n",
|
||||||
|
" if self.modulate:\n",
|
||||||
|
" decay = torch.exp(-t * self.deltas.abs()) \n",
|
||||||
|
" x = x * (decay + self.shift)\n",
|
||||||
|
" return x \n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"class HyenaFilter(OptimModule):\n",
|
||||||
|
" def __init__(\n",
|
||||||
|
" self, \n",
|
||||||
|
" d_model,\n",
|
||||||
|
" emb_dim=3, # dim of input to MLP, augments with positional encoding\n",
|
||||||
|
" order=16, # width of the implicit MLP \n",
|
||||||
|
" fused_fft_conv=False,\n",
|
||||||
|
" seq_len=1024, \n",
|
||||||
|
" lr=1e-3, \n",
|
||||||
|
" lr_pos_emb=1e-5,\n",
|
||||||
|
" dropout=0.0, \n",
|
||||||
|
" w=1, # frequency of periodic activations \n",
|
||||||
|
" wd=0, # weight decay of kernel parameters \n",
|
||||||
|
" bias=True,\n",
|
||||||
|
" num_inner_mlps=2,\n",
|
||||||
|
" normalized=False,\n",
|
||||||
|
" **kwargs\n",
|
||||||
|
" ):\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" Implicit long filter with modulation.\n",
|
||||||
|
" \n",
|
||||||
|
" Args:\n",
|
||||||
|
" d_model: number of channels in the input\n",
|
||||||
|
" emb_dim: dimension of the positional encoding (`emb_dim` - 1) // 2 is the number of bands\n",
|
||||||
|
" order: width of the FFN\n",
|
||||||
|
" num_inner_mlps: number of inner linear layers inside filter MLP\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" self.d_model = d_model\n",
|
||||||
|
" self.use_bias = bias\n",
|
||||||
|
" self.fused_fft_conv = fused_fft_conv\n",
|
||||||
|
" self.bias = nn.Parameter(torch.randn(self.d_model))\n",
|
||||||
|
" self.dropout = nn.Dropout(dropout)\n",
|
||||||
|
" \n",
|
||||||
|
" act = Sin(dim=order, w=w)\n",
|
||||||
|
" self.emb_dim = emb_dim\n",
|
||||||
|
" assert emb_dim % 2 != 0 and emb_dim >= 3, \"emb_dim must be odd and greater or equal to 3 (time, sine and cosine)\"\n",
|
||||||
|
" self.seq_len = seq_len\n",
|
||||||
|
" \n",
|
||||||
|
" self.pos_emb = PositionalEmbedding(emb_dim, seq_len, lr_pos_emb)\n",
|
||||||
|
" \n",
|
||||||
|
" self.implicit_filter = nn.Sequential(\n",
|
||||||
|
" nn.Linear(emb_dim, order),\n",
|
||||||
|
" act,\n",
|
||||||
|
" )\n",
|
||||||
|
" for i in range(num_inner_mlps):\n",
|
||||||
|
" self.implicit_filter.append(nn.Linear(order, order))\n",
|
||||||
|
" self.implicit_filter.append(act)\n",
|
||||||
|
"\n",
|
||||||
|
" self.implicit_filter.append(nn.Linear(order, d_model, bias=False))\n",
|
||||||
|
" \n",
|
||||||
|
" self.modulation = ExponentialModulation(d_model, **kwargs)\n",
|
||||||
|
" \n",
|
||||||
|
" self.normalized = normalized\n",
|
||||||
|
" for c in self.implicit_filter.children():\n",
|
||||||
|
" for name, v in c.state_dict().items(): \n",
|
||||||
|
" optim = {\"weight_decay\": wd, \"lr\": lr}\n",
|
||||||
|
" setattr(getattr(c, name), \"_optim\", optim)\n",
|
||||||
|
"\n",
|
||||||
|
" def filter(self, L, *args, **kwargs):\n",
|
||||||
|
" z, t = self.pos_emb(L)\n",
|
||||||
|
" h = self.implicit_filter(z)\n",
|
||||||
|
" h = self.modulation(t, h)\n",
|
||||||
|
" return h\n",
|
||||||
|
"\n",
|
||||||
|
" def forward(self, x, L, k=None, bias=None, *args, **kwargs):\n",
|
||||||
|
" if k is None: k = self.filter(L)\n",
|
||||||
|
" \n",
|
||||||
|
" # Ensure compatibility with filters that return a tuple \n",
|
||||||
|
" k = k[0] if type(k) is tuple else k \n",
|
||||||
|
"\n",
|
||||||
|
" y = fftconv(x, k, bias)\n",
|
||||||
|
" return y\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
"class HyenaOperator(nn.Module):\n",
|
||||||
|
" def __init__(\n",
|
||||||
|
" self,\n",
|
||||||
|
" d_model,\n",
|
||||||
|
" l_max,\n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64,\n",
|
||||||
|
" dropout=0.0, \n",
|
||||||
|
" filter_dropout=0.0, \n",
|
||||||
|
" **filter_args,\n",
|
||||||
|
" ):\n",
|
||||||
|
" r\"\"\"\n",
|
||||||
|
" Hyena operator described in the paper https://arxiv.org/pdf/2302.10866.pdf\n",
|
||||||
|
" \n",
|
||||||
|
" Args:\n",
|
||||||
|
" d_model (int): Dimension of the input and output embeddings (width of the layer)\n",
|
||||||
|
" l_max: (int): Maximum input sequence length. Defaults to None\n",
|
||||||
|
" order: (int): Depth of the Hyena recurrence. Defaults to 2\n",
|
||||||
|
" dropout: (float): Dropout probability. Defaults to 0.0\n",
|
||||||
|
" filter_dropout: (float): Dropout probability for the filter. Defaults to 0.0\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" self.d_model = d_model\n",
|
||||||
|
" self.l_max = l_max\n",
|
||||||
|
" self.order = order\n",
|
||||||
|
" inner_width = d_model * (order + 1)\n",
|
||||||
|
" self.dropout = nn.Dropout(dropout)\n",
|
||||||
|
" self.in_proj = nn.Linear(d_model, inner_width)\n",
|
||||||
|
" self.out_proj = nn.Linear(d_model, d_model)\n",
|
||||||
|
" \n",
|
||||||
|
" self.short_filter = nn.Conv1d(\n",
|
||||||
|
" inner_width, \n",
|
||||||
|
" inner_width, \n",
|
||||||
|
" 3,\n",
|
||||||
|
" padding=2,\n",
|
||||||
|
" groups=inner_width\n",
|
||||||
|
" )\n",
|
||||||
|
" self.filter_fn = HyenaFilter(\n",
|
||||||
|
" d_model * (order - 1), \n",
|
||||||
|
" order=filter_order, \n",
|
||||||
|
" seq_len=l_max,\n",
|
||||||
|
" channels=1, \n",
|
||||||
|
" dropout=filter_dropout, \n",
|
||||||
|
" **filter_args\n",
|
||||||
|
" ) \n",
|
||||||
|
"\n",
|
||||||
|
" def forward(self, u, *args, **kwargs):\n",
|
||||||
|
" l = u.size(-2)\n",
|
||||||
|
" l_filter = min(l, self.l_max)\n",
|
||||||
|
" u = self.in_proj(u)\n",
|
||||||
|
" u = rearrange(u, 'b l d -> b d l')\n",
|
||||||
|
" \n",
|
||||||
|
" uc = self.short_filter(u)[...,:l_filter] \n",
|
||||||
|
" *x, v = uc.split(self.d_model, dim=1)\n",
|
||||||
|
" \n",
|
||||||
|
" k = self.filter_fn.filter(l_filter)[0]\n",
|
||||||
|
" k = rearrange(k, 'l (o d) -> o d l', o=self.order - 1)\n",
|
||||||
|
" bias = rearrange(self.filter_fn.bias, '(o d) -> o d', o=self.order - 1)\n",
|
||||||
|
" \n",
|
||||||
|
" for o, x_i in enumerate(reversed(x[1:])):\n",
|
||||||
|
" v = self.dropout(v * x_i)\n",
|
||||||
|
" v = self.filter_fn(v, l_filter, k=k[o], bias=bias[o])\n",
|
||||||
|
"\n",
|
||||||
|
" y = rearrange(v * x[0], 'b d l -> b l d')\n",
|
||||||
|
"\n",
|
||||||
|
" y = self.out_proj(y)\n",
|
||||||
|
" return y\n",
|
||||||
|
"\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
"if __name__ == \"__main__\":\n",
|
||||||
|
" layer = HyenaOperator(\n",
|
||||||
|
" \n",
|
||||||
|
" d_model=128, \n",
|
||||||
|
" l_max=1024, \n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64\n",
|
||||||
|
" )\n",
|
||||||
|
" x = torch.randn(1, 1024, 128, requires_grad=True)\n",
|
||||||
|
" y = layer(x)\n",
|
||||||
|
" \n",
|
||||||
|
" print(x.shape, y.shape)\n",
|
||||||
|
" \n",
|
||||||
|
" grad = torch.autograd.grad(y[:, 10, :].sum(), x)[0]\n",
|
||||||
|
" print('Causality check: gradients should not flow \"from future to past\"')\n",
|
||||||
|
" print(grad[0, 11, :].sum(), grad[0, 9, :].sum())\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 368,
|
||||||
|
"id": "032ef08a-8cc6-491a-9eb8-4a6b3f2d165e",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"torch.Size([1, 1023, 1]) torch.Size([1, 500])\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"class HyenaOperatorAutoregressive1D(nn.Module):\n",
|
||||||
|
" def __init__(\n",
|
||||||
|
" self,\n",
|
||||||
|
" d_model,\n",
|
||||||
|
" l_max,\n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64,\n",
|
||||||
|
" dropout=0.0, \n",
|
||||||
|
" filter_dropout=0.0, \n",
|
||||||
|
" **filter_args,\n",
|
||||||
|
" ):\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
"\n",
|
||||||
|
" self.l_max = l_max\n",
|
||||||
|
" self.d_model = d_model\n",
|
||||||
|
" self.l_max = l_max\n",
|
||||||
|
" self.order = order\n",
|
||||||
|
" inner_width = d_model * (order + 1)\n",
|
||||||
|
"\n",
|
||||||
|
" self.dropout = nn.Dropout(dropout)\n",
|
||||||
|
" self.in_proj = nn.Linear(d_model, inner_width)\n",
|
||||||
|
" self.out_proj = nn.Linear(d_model, d_model)\n",
|
||||||
|
" self.fc_before = nn.Linear(1, d_model) # Fully connected layer before the main layer\n",
|
||||||
|
" self.fc_after = nn.Linear(d_model, 500) # Fully connected layer after the main layer\n",
|
||||||
|
"\n",
|
||||||
|
" self.operator = HyenaOperator(\n",
|
||||||
|
" d_model=d_model,\n",
|
||||||
|
" l_max=l_max,\n",
|
||||||
|
" order=order, \n",
|
||||||
|
" filter_order=filter_order,\n",
|
||||||
|
" dropout=dropout, \n",
|
||||||
|
" filter_dropout=filter_dropout, \n",
|
||||||
|
" **filter_args,\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" def forward(self, u, *args, **kwargs):\n",
|
||||||
|
" # Increase the channel dimension from 1 to d_model\n",
|
||||||
|
" u = self.fc_before(u) \n",
|
||||||
|
" # Pass through the operator\n",
|
||||||
|
" #[B,1024,128] --> [B,1024,128]\n",
|
||||||
|
" u = self.operator(u)\n",
|
||||||
|
" \n",
|
||||||
|
" last_state = u[:,-1,:]\n",
|
||||||
|
" # Decrease the channel dimension back to 1\n",
|
||||||
|
" y = self.fc_after(last_state)\n",
|
||||||
|
" return y,last_state\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"if __name__ == \"__main__\":\n",
|
||||||
|
" layer = HyenaOperatorAutoregressive1D(\n",
|
||||||
|
" d_model=128, \n",
|
||||||
|
" l_max=1024, \n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" x = torch.randn(1, 1023, 1, requires_grad=True) # 1D time series input\n",
|
||||||
|
" y, last_state = layer(x)\n",
|
||||||
|
"\n",
|
||||||
|
" #import pdb;pdb.set_trace()\n",
|
||||||
|
" print(x.shape, y.shape) # should now be [1, 1024, 1]\n",
|
||||||
|
"\n",
|
||||||
|
" #grad = torch.autograd.grad(y[:, 10, 0].sum(), x)[0]\n",
|
||||||
|
" #print('Causality check: gradients should not flow \"from future to past\"')\n",
|
||||||
|
" #print(grad[0, 11, 0].sum(), grad[0, 9, 0].sum())"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 369,
|
||||||
|
"id": "80cde67b-992f-4cb0-8824-4a6b7e4984ca",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"Train Epoch: 1 [0/640 (0%)]\tLoss: 0.569402\n",
|
||||||
|
"Train Epoch: 2 [0/640 (0%)]\tLoss: 0.459675\n",
|
||||||
|
"Train Epoch: 3 [0/640 (0%)]\tLoss: 0.455738\n",
|
||||||
|
"Train Epoch: 4 [0/640 (0%)]\tLoss: 0.334769\n",
|
||||||
|
"Train Epoch: 5 [0/640 (0%)]\tLoss: 0.400367\n",
|
||||||
|
"Train Epoch: 6 [0/640 (0%)]\tLoss: 0.353075\n",
|
||||||
|
"Train Epoch: 7 [0/640 (0%)]\tLoss: 0.372000\n",
|
||||||
|
"Train Epoch: 8 [0/640 (0%)]\tLoss: 0.354896\n",
|
||||||
|
"Train Epoch: 9 [0/640 (0%)]\tLoss: 0.370209\n",
|
||||||
|
"Train Epoch: 10 [0/640 (0%)]\tLoss: 0.323172\n",
|
||||||
|
"Train Epoch: 11 [0/640 (0%)]\tLoss: 0.324925\n",
|
||||||
|
"Train Epoch: 12 [0/640 (0%)]\tLoss: 0.319851\n",
|
||||||
|
"Train Epoch: 13 [0/640 (0%)]\tLoss: 0.267247\n",
|
||||||
|
"Train Epoch: 14 [0/640 (0%)]\tLoss: 0.307432\n",
|
||||||
|
"Train Epoch: 15 [0/640 (0%)]\tLoss: 0.270438\n",
|
||||||
|
"Train Epoch: 16 [0/640 (0%)]\tLoss: 0.274847\n",
|
||||||
|
"Train Epoch: 17 [0/640 (0%)]\tLoss: 0.283570\n",
|
||||||
|
"Train Epoch: 18 [0/640 (0%)]\tLoss: 0.282435\n",
|
||||||
|
"Train Epoch: 19 [0/640 (0%)]\tLoss: 0.325277\n",
|
||||||
|
"Train Epoch: 20 [0/640 (0%)]\tLoss: 0.326561\n",
|
||||||
|
"Train Epoch: 21 [0/640 (0%)]\tLoss: 0.305762\n",
|
||||||
|
"Train Epoch: 22 [0/640 (0%)]\tLoss: 0.298317\n",
|
||||||
|
"Train Epoch: 23 [0/640 (0%)]\tLoss: 0.277244\n",
|
||||||
|
"Train Epoch: 24 [0/640 (0%)]\tLoss: 0.271200\n",
|
||||||
|
"Train Epoch: 25 [0/640 (0%)]\tLoss: 0.246922\n",
|
||||||
|
"Train Epoch: 26 [0/640 (0%)]\tLoss: 0.231166\n",
|
||||||
|
"Train Epoch: 27 [0/640 (0%)]\tLoss: 0.301950\n",
|
||||||
|
"Train Epoch: 28 [0/640 (0%)]\tLoss: 0.272427\n",
|
||||||
|
"Train Epoch: 29 [0/640 (0%)]\tLoss: 0.278342\n",
|
||||||
|
"Train Epoch: 30 [0/640 (0%)]\tLoss: 0.246955\n",
|
||||||
|
"Train Epoch: 31 [0/640 (0%)]\tLoss: 0.304109\n",
|
||||||
|
"Train Epoch: 32 [0/640 (0%)]\tLoss: 0.254574\n",
|
||||||
|
"Train Epoch: 33 [0/640 (0%)]\tLoss: 0.234942\n",
|
||||||
|
"Train Epoch: 34 [0/640 (0%)]\tLoss: 0.269669\n",
|
||||||
|
"Train Epoch: 35 [0/640 (0%)]\tLoss: 0.239477\n",
|
||||||
|
"Train Epoch: 36 [0/640 (0%)]\tLoss: 0.282838\n",
|
||||||
|
"Train Epoch: 37 [0/640 (0%)]\tLoss: 0.264821\n",
|
||||||
|
"Train Epoch: 38 [0/640 (0%)]\tLoss: 0.236031\n",
|
||||||
|
"Train Epoch: 39 [0/640 (0%)]\tLoss: 0.239743\n",
|
||||||
|
"Train Epoch: 40 [0/640 (0%)]\tLoss: 0.289715\n",
|
||||||
|
"Train Epoch: 41 [0/640 (0%)]\tLoss: 0.240897\n",
|
||||||
|
"Train Epoch: 42 [0/640 (0%)]\tLoss: 0.212328\n",
|
||||||
|
"Train Epoch: 43 [0/640 (0%)]\tLoss: 0.253719\n",
|
||||||
|
"Train Epoch: 44 [0/640 (0%)]\tLoss: 0.210382\n",
|
||||||
|
"Train Epoch: 45 [0/640 (0%)]\tLoss: 0.284465\n",
|
||||||
|
"Train Epoch: 46 [0/640 (0%)]\tLoss: 0.292590\n",
|
||||||
|
"Train Epoch: 47 [0/640 (0%)]\tLoss: 0.200158\n",
|
||||||
|
"Train Epoch: 48 [0/640 (0%)]\tLoss: 0.269204\n",
|
||||||
|
"Train Epoch: 49 [0/640 (0%)]\tLoss: 0.213730\n",
|
||||||
|
"Train Epoch: 50 [0/640 (0%)]\tLoss: 0.268548\n",
|
||||||
|
"Train Epoch: 51 [0/640 (0%)]\tLoss: 0.186853\n",
|
||||||
|
"Train Epoch: 52 [0/640 (0%)]\tLoss: 0.263178\n",
|
||||||
|
"Train Epoch: 53 [0/640 (0%)]\tLoss: 0.232999\n",
|
||||||
|
"Train Epoch: 54 [0/640 (0%)]\tLoss: 0.265539\n",
|
||||||
|
"Train Epoch: 55 [0/640 (0%)]\tLoss: 0.227278\n",
|
||||||
|
"Train Epoch: 56 [0/640 (0%)]\tLoss: 0.266715\n",
|
||||||
|
"Train Epoch: 57 [0/640 (0%)]\tLoss: 0.243662\n",
|
||||||
|
"Train Epoch: 58 [0/640 (0%)]\tLoss: 0.231188\n",
|
||||||
|
"Train Epoch: 59 [0/640 (0%)]\tLoss: 0.234178\n",
|
||||||
|
"Train Epoch: 60 [0/640 (0%)]\tLoss: 0.270513\n",
|
||||||
|
"Train Epoch: 61 [0/640 (0%)]\tLoss: 0.240779\n",
|
||||||
|
"Train Epoch: 62 [0/640 (0%)]\tLoss: 0.278103\n",
|
||||||
|
"Train Epoch: 63 [0/640 (0%)]\tLoss: 0.235614\n",
|
||||||
|
"Train Epoch: 64 [0/640 (0%)]\tLoss: 0.248226\n",
|
||||||
|
"Train Epoch: 65 [0/640 (0%)]\tLoss: 0.242527\n",
|
||||||
|
"Train Epoch: 66 [0/640 (0%)]\tLoss: 0.312065\n",
|
||||||
|
"Train Epoch: 67 [0/640 (0%)]\tLoss: 0.201206\n",
|
||||||
|
"Train Epoch: 68 [0/640 (0%)]\tLoss: 0.263827\n",
|
||||||
|
"Train Epoch: 69 [0/640 (0%)]\tLoss: 0.286222\n",
|
||||||
|
"Train Epoch: 70 [0/640 (0%)]\tLoss: 0.279344\n",
|
||||||
|
"Train Epoch: 71 [0/640 (0%)]\tLoss: 0.263106\n",
|
||||||
|
"Train Epoch: 72 [0/640 (0%)]\tLoss: 0.247235\n",
|
||||||
|
"Train Epoch: 73 [0/640 (0%)]\tLoss: 0.243001\n",
|
||||||
|
"Train Epoch: 74 [0/640 (0%)]\tLoss: 0.231101\n",
|
||||||
|
"Train Epoch: 75 [0/640 (0%)]\tLoss: 0.290737\n",
|
||||||
|
"Train Epoch: 76 [0/640 (0%)]\tLoss: 0.274286\n",
|
||||||
|
"Train Epoch: 77 [0/640 (0%)]\tLoss: 0.254139\n",
|
||||||
|
"Train Epoch: 78 [0/640 (0%)]\tLoss: 0.296462\n",
|
||||||
|
"Train Epoch: 79 [0/640 (0%)]\tLoss: 0.223839\n",
|
||||||
|
"Train Epoch: 80 [0/640 (0%)]\tLoss: 0.282157\n",
|
||||||
|
"Train Epoch: 81 [0/640 (0%)]\tLoss: 0.193232\n",
|
||||||
|
"Train Epoch: 82 [0/640 (0%)]\tLoss: 0.319076\n",
|
||||||
|
"Train Epoch: 83 [0/640 (0%)]\tLoss: 0.195717\n",
|
||||||
|
"Train Epoch: 84 [0/640 (0%)]\tLoss: 0.243448\n",
|
||||||
|
"Train Epoch: 85 [0/640 (0%)]\tLoss: 0.218913\n",
|
||||||
|
"Train Epoch: 86 [0/640 (0%)]\tLoss: 0.267013\n",
|
||||||
|
"Train Epoch: 87 [0/640 (0%)]\tLoss: 0.243178\n",
|
||||||
|
"Train Epoch: 88 [0/640 (0%)]\tLoss: 0.211878\n",
|
||||||
|
"Train Epoch: 89 [0/640 (0%)]\tLoss: 0.294469\n",
|
||||||
|
"Train Epoch: 90 [0/640 (0%)]\tLoss: 0.270843\n",
|
||||||
|
"Train Epoch: 91 [0/640 (0%)]\tLoss: 0.212928\n",
|
||||||
|
"Train Epoch: 92 [0/640 (0%)]\tLoss: 0.284125\n",
|
||||||
|
"Train Epoch: 93 [0/640 (0%)]\tLoss: 0.236006\n",
|
||||||
|
"Train Epoch: 94 [0/640 (0%)]\tLoss: 0.260203\n",
|
||||||
|
"Train Epoch: 95 [0/640 (0%)]\tLoss: 0.270325\n",
|
||||||
|
"Train Epoch: 96 [0/640 (0%)]\tLoss: 0.199420\n",
|
||||||
|
"Train Epoch: 97 [0/640 (0%)]\tLoss: 0.313867\n",
|
||||||
|
"Train Epoch: 98 [0/640 (0%)]\tLoss: 0.285310\n",
|
||||||
|
"Train Epoch: 99 [0/640 (0%)]\tLoss: 0.284724\n",
|
||||||
|
"Train Epoch: 100 [0/640 (0%)]\tLoss: 0.256413\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"import torch\n",
|
||||||
|
"import torch.optim as optim\n",
|
||||||
|
"import torch.nn.functional as F\n",
|
||||||
|
"from torch.utils.data import DataLoader, Dataset\n",
|
||||||
|
"import numpy as np\n",
|
||||||
|
"\n",
|
||||||
|
"def generate_sine_with_noise(n_points, frequency, phase, amplitude, noise_sd):\n",
|
||||||
|
" # Generate an array of points from 0 to 2*pi\n",
|
||||||
|
" x = np.linspace(0, 2*np.pi, n_points)\n",
|
||||||
|
" \n",
|
||||||
|
" # Generate the sine wave\n",
|
||||||
|
" sine_wave = amplitude * np.sin(frequency * x + phase)\n",
|
||||||
|
" \n",
|
||||||
|
" # Generate Gaussian noise\n",
|
||||||
|
" noise = np.random.normal(scale=noise_sd, size=n_points)\n",
|
||||||
|
" \n",
|
||||||
|
" # Add the noise to the sine wave\n",
|
||||||
|
" sine_wave_noise = sine_wave + noise\n",
|
||||||
|
" \n",
|
||||||
|
" # Stack the sine wave and the noisy sine wave into a 2D array\n",
|
||||||
|
" output = np.column_stack((sine_wave, sine_wave_noise))\n",
|
||||||
|
" \n",
|
||||||
|
" return output\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
"class SineDataset(Dataset):\n",
|
||||||
|
" def __init__(self, n_samples, n_points, frequency_range, phase_range, amplitude_range, noise_sd_range):\n",
|
||||||
|
" self.n_samples = n_samples\n",
|
||||||
|
" self.n_points = n_points\n",
|
||||||
|
" self.frequency_range = frequency_range\n",
|
||||||
|
" self.phase_range = phase_range\n",
|
||||||
|
" self.amplitude_range = amplitude_range\n",
|
||||||
|
" self.noise_sd_range = noise_sd_range\n",
|
||||||
|
"\n",
|
||||||
|
" def __len__(self):\n",
|
||||||
|
" return self.n_samples\n",
|
||||||
|
"\n",
|
||||||
|
" def __getitem__(self, idx):\n",
|
||||||
|
" # Generate random attributes\n",
|
||||||
|
" frequency = np.random.uniform(*self.frequency_range)\n",
|
||||||
|
" phase = np.random.uniform(*self.phase_range)\n",
|
||||||
|
" amplitude = np.random.uniform(*self.amplitude_range)\n",
|
||||||
|
" noise_sd = np.random.uniform(*self.noise_sd_range)\n",
|
||||||
|
"\n",
|
||||||
|
" # Generate sine wave with the random attributes\n",
|
||||||
|
" sine_wave = generate_sine_with_noise(self.n_points, frequency, phase, amplitude, noise_sd)\n",
|
||||||
|
"\n",
|
||||||
|
" # Return the sine wave and the parameters\n",
|
||||||
|
" return torch.Tensor(sine_wave[:-500, 1, None]), torch.Tensor(sine_wave[-500:, 0]), torch.Tensor([frequency, phase, amplitude, noise_sd])\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# Usage:\n",
|
||||||
|
"dataset = SineDataset(640, 1025, (1, 3), (0, 2*np.pi), (0.5, 1.5), (0.05, 1))\n",
|
||||||
|
"\n",
|
||||||
|
"def train(model, device, train_loader, optimizer, epoch):\n",
|
||||||
|
" model.train()\n",
|
||||||
|
" for batch_idx, (data, target, params) in enumerate(train_loader):\n",
|
||||||
|
" #data = data[...,None]\n",
|
||||||
|
" data, target = data.to(device), target.to(device)\n",
|
||||||
|
" optimizer.zero_grad()\n",
|
||||||
|
" output,last_state = model(data)\n",
|
||||||
|
" #import pdb;pdb.set_trace()\n",
|
||||||
|
"\n",
|
||||||
|
" loss = F.mse_loss(output, target)\n",
|
||||||
|
" loss.backward()\n",
|
||||||
|
" optimizer.step()\n",
|
||||||
|
" if batch_idx % 10 == 0:\n",
|
||||||
|
" print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
|
||||||
|
" epoch, batch_idx * len(data), len(train_loader.dataset),\n",
|
||||||
|
" 100. * batch_idx / len(train_loader), loss.item()))\n",
|
||||||
|
"\n",
|
||||||
|
"if __name__ == \"__main__\":\n",
|
||||||
|
" device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
|
||||||
|
"\n",
|
||||||
|
" model = HyenaOperatorAutoregressive1D(\n",
|
||||||
|
" d_model=128, \n",
|
||||||
|
" l_max=1024, \n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64\n",
|
||||||
|
" ).to(device)\n",
|
||||||
|
"\n",
|
||||||
|
" optimizer = optim.Adam(model.parameters())\n",
|
||||||
|
"\n",
|
||||||
|
" # Assume 10000 samples in the dataset\n",
|
||||||
|
" #dataset = SineDataset(10000, 1025, 2, 0, 1, 0.1)\n",
|
||||||
|
" train_loader = DataLoader(dataset, batch_size=64, shuffle=True)\n",
|
||||||
|
"\n",
|
||||||
|
" for epoch in range(1, 101): # Train for 10 epochs\n",
|
||||||
|
" train(model, device, train_loader, optimizer, epoch)\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "cc9f9031-5ee1-49f8-a70f-ad85ca015596",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 682,
|
||||||
|
"id": "90330622-8b44-4b45-8158-6840538f768c",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"from sklearn.linear_model import LinearRegression\n",
|
||||||
|
"from sklearn.metrics import r2_score\n",
|
||||||
|
"\n",
|
||||||
|
"def fit_and_evaluate_linear_regression(outputs_and_params):\n",
|
||||||
|
"\n",
|
||||||
|
" # Split the data into inputs (last_states) and targets (params)\n",
|
||||||
|
"\n",
|
||||||
|
" inputs = np.concatenate([x[0] for x in outputs_and_params])\n",
|
||||||
|
"\n",
|
||||||
|
" targets = np.concatenate([x[1] for x in outputs_and_params])\n",
|
||||||
|
"\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
" r2_scores = []\n",
|
||||||
|
"\n",
|
||||||
|
" param_names = [\"frequency\", \"sin_phase\", \"amplitude\", \"noise_sd\"]\n",
|
||||||
|
"\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
" # Fit the linear regression model for each parameter and calculate the R^2 score\n",
|
||||||
|
"\n",
|
||||||
|
" for i in range(targets.shape[1]):\n",
|
||||||
|
"\n",
|
||||||
|
" if param_names[i] == 'sin_phase':\n",
|
||||||
|
" # take the sine of the phase values\n",
|
||||||
|
" target_values = np.sin(targets[:, i])\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
" #Not used at the moment\n",
|
||||||
|
" elif param_names[i] == 'cos_phase':\n",
|
||||||
|
" # take the cosine of the phase values\n",
|
||||||
|
" target_values = np.cos(targets[:, i])\n",
|
||||||
|
" else:\n",
|
||||||
|
" target_values = targets[:, i]\n",
|
||||||
|
"\n",
|
||||||
|
" model = LinearRegression().fit(inputs, target_values)\n",
|
||||||
|
"\n",
|
||||||
|
" pred = model.predict(inputs)\n",
|
||||||
|
"\n",
|
||||||
|
" score = r2_score(target_values, pred)\n",
|
||||||
|
"\n",
|
||||||
|
" r2_scores.append(score)\n",
|
||||||
|
"\n",
|
||||||
|
" print(f\"R^2 score for {param_names[i]}: {score:.2f}\")\n",
|
||||||
|
"\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
" return r2_scores"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 683,
|
||||||
|
"id": "5eb62a22-cad8-43c4-b757-f36b6a01e9be",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def generate_outputs(model, device, data_loader):\n",
|
||||||
|
" model.eval()\n",
|
||||||
|
" outputs_and_params = []\n",
|
||||||
|
" with torch.no_grad():\n",
|
||||||
|
" for data, target, params in data_loader:\n",
|
||||||
|
" data, target = data.to(device), target.to(device)\n",
|
||||||
|
" output, last_state = model(data)\n",
|
||||||
|
" outputs_and_params.append((last_state.cpu().numpy(), params.cpu().numpy()))\n",
|
||||||
|
" return outputs_and_params\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 684,
|
||||||
|
"id": "a95ee542-1c39-4f04-9184-e26c6983a018",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"R^2 score for frequency: -1.09\n",
|
||||||
|
"R^2 score for sin_phase: 0.84\n",
|
||||||
|
"R^2 score for amplitude: 0.96\n",
|
||||||
|
"R^2 score for noise_sd: 0.91\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"outputs_and_params = generate_outputs(model, device, train_loader)\n",
|
||||||
|
"\n",
|
||||||
|
"# Fit the linear regression model and print the R^2 score for each parameter\n",
|
||||||
|
"r2_scores = fit_and_evaluate_linear_regression(outputs_and_params)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "de65d0d2-b0c6-4ac5-a87f-70fa7f90480b",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"# Autoregressive"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 519,
|
||||||
|
"id": "8ee139a6-aee5-4309-8685-bf0c28893279",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def predict_autoregressive(model, initial_data, n_steps):\n",
|
||||||
|
" model.eval()\n",
|
||||||
|
" predictions = []\n",
|
||||||
|
" current_input = initial_data\n",
|
||||||
|
" with torch.no_grad():\n",
|
||||||
|
" for _ in range(n_steps):\n",
|
||||||
|
" # Get the prediction for the next step and save it\n",
|
||||||
|
" next_output, last_state = model(current_input)\n",
|
||||||
|
" predictions.append(next_output)\n",
|
||||||
|
" #import pdb;pdb.set_trace()\n",
|
||||||
|
"\n",
|
||||||
|
" # Prepare the input for the next step\n",
|
||||||
|
" next_input = torch.cat((current_input[:, 500:, :], next_output[:,:,None]), dim=1)\n",
|
||||||
|
"\n",
|
||||||
|
" current_input = next_input\n",
|
||||||
|
"\n",
|
||||||
|
" return torch.cat(predictions, dim=1)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "272040d6-4dfb-438b-a432-744e950effd1",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 594,
|
||||||
|
"id": "9f49cbf8-08e8-428a-af80-bcb2515a6327",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"torch.Size([1, 525, 1])"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 594,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"initial_data = dataset[0]\n",
|
||||||
|
"initial_data[0][None,...].shape"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 595,
|
||||||
|
"id": "c75464e3-e44d-4325-8804-29d8081e3a45",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"autoregressive_out = predict_autoregressive(model,initial_data[0][None,...],3)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 596,
|
||||||
|
"id": "90a37c56-59f3-49bc-bfbe-d9e126c42ed1",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"total = np.concatenate([initial_data[0].squeeze().numpy() ,autoregressive_out.squeeze().numpy()])"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 597,
|
||||||
|
"id": "3d2ad0ee-7a0e-4b6f-8b40-dd565c8af84d",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"[<matplotlib.lines.Line2D at 0x7f8ec058feb0>]"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 597,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAiIAAAGhCAYAAABceN/BAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAABehElEQVR4nO3dd3xUVfo/8M+kTXoCabRA6EgPCCGoSFNA1l5YRAWWxbKwFpQV3FXssKtf3dVVZP0JuCp2hVUUpCPSS+idhARIAgTSCKlzf39M5uZOMjW5debzfr2id2Zu7j0zE+Y+c85znmMSBEEAERERkQYCtG4AERER+S8GIkRERKQZBiJERESkGQYiREREpBkGIkRERKQZBiJERESkGQYiREREpBkGIkRERKQZBiJERESkGQYiREREpBnVApF58+bBZDLhySefVOuUREREpHOqBCI7duzAggUL0Lt3bzVOR0RERAYRpPQJSktLMWHCBHz44Yd49dVXvfpdi8WCc+fOISoqCiaTSaEWEhERkZwEQUBJSQlatWqFgADXfR6KByLTpk3D2LFjMXLkSLeBSEVFBSoqKsTbZ8+eRffu3ZVuIhERESkgJycHbdq0cbmPooHIF198gd27d2PHjh0e7T937ly89NJLDe7PyclBdHS03M0jIiIiBRQXFyM5ORlRUVFu91UsEMnJycETTzyBVatWITQ01KPfmT17NmbMmCHetj2R6OhoBiJEREQG40lahUkQBEGJky9duhR33nknAgMDxftqampgMpkQEBCAiooKu8ccKS4uRkxMDIqKihiIEBERGYQ312/FekRGjBiB/fv32903efJkdOvWDc8++6zbIISIiIh8n2KBSFRUFHr27Gl3X0REBOLi4hrcT0RERP6JlVWJiIhIM4pP35Vav369mqcjIiIinWOPCBEREWmGgQgRERFphoEIERERaYaBCBEREWmGgQgRERFphoEIERERaYaBCBEREWmGgYif2XaqAJ9vz9a6GURERABULmhG2hv3n60AgI4JkRjYvrnGrSEiIn/HHhE/lX2pTOsmEBERMRAhIiIi7TAQISIiIs0wEPFTJq0bQEREBAYiREREpCEGIkRERKQZTt/1EzmXyhAeEqh1M4iIiOwwEPEDBaUVuOEf6+zuMzFJhIiIdIBDM37gaH6J1k0gIiJyiIEIERERaYaBCBEREWmGgQgRERFphoEIERERaYaBiJ/irBkiItIDBiJERESkGQYifsDElWWIiEinGIj4KQYnRESkBwxEiIiISDMMRIiIiEgzDESIiIhIMwxE/ICjqbqcvktERHrAQISIiIg0w0CEiIiINMNAhIiIiDTDQISIiIg0w0CEiIiINMNAhIiIiDSjaCAyf/589O7dG9HR0YiOjkZ6ejp+/vlnJU9JDuzMuqR1E4iIiBxSNBBp06YN5s2bh127dmHnzp0YPnw4br/9dhw8eFDJ01I9b/5yTOsmEBERORSk5MFvvfVWu9uvvfYa5s+fj61bt6JHjx5KnpqIiIgMQNFARKqmpgZff/01rly5gvT0dLVOS06YWFqViIh0QPFAZP/+/UhPT0d5eTkiIyPx/fffo3v37g73raioQEVFhXi7uLhY6eYRERGRhhSfNdO1a1dkZGRg27ZteOyxxzBx4kQcOnTI4b5z585FTEyM+JOcnKx084iIiEhDJkEQBDVPOHLkSHTs2BELFixo8JijHpHk5GQUFRUhOjpazWb6lJRZyxvc9874VNzWp5UGrSEiIl9XXFyMmJgYj67fquWI2FgsFrtgQ8psNsNsNqvcIv/EDBEiItIDRQOR2bNnY8yYMWjbti1KSkqwZMkSrF+/HitXrlTytCSzGouAi6UVSIoO1bopRETkYxQNRM6fP4+HHnoIubm5iImJQe/evbFy5UrcdNNNSp6WZDZx4XZsOnERn08dhPSOcVo3h4iIfIiigchHH32k5OFJJZtOXAQAfLrtNAMRIiKSFdea8WNf7czBwk2ZWjeDiIj8mOrJqqQff/lmHwBgVM8WaB0bpnFriIjIH7FHxE9J52xfqajWrB1EROTfGIgQERGRZhiI+CmV69gRERE5xECEWNyMiIg0w0DETz2/9IDXv8OAhYiI5MZAxE8VlzNBlYiItMdAhGBiVwcREWmEgQgRERFphoEIERERaYaBiI+yWATkFZV7uDfHZoiISBsMRHzU41/swaC5a7DyYJ7WTSEiInKKgYiP+nFfLgDgkU92adwSIiIi5xiIkMezZkycXkNERDJjIEJERESaYSDiQ9YdOY+cS2VaN4OIiMhjQVo3gOSx4dgFTF68AwCQNW+sxq0hIiLyDHtEfMTOrEuN/l1mfhARkVYYiBAREZFmGIiQx9hzQkREcmMgQkRERJphIEJERESaYSBCLFRGRESaYSDiIxhKEBGRETEQISIiIs0wECGHvSl7cwox8+u9OF9Srnp7iIjIf7CyKjlc9O72934DAFworVC5NURE5E/YI0IunThfKm4zp5WIiOTGQIRcEgStW0BERL6MgYgPen/9Ca/2twjA0bwSWCyMOoiISF0MRHzQP1Yc9Wr/uT8dxqh/bsTfVx5p8JjALhEiIlIQAxHCL4fyAQALNpzSuCVERORvGIgQERGRZhiIkEscmCEiIiUxEPEVnFtLREQGpGggMnfuXAwYMABRUVFITEzEHXfcgaNHvUukJG1Jc1UZ6hARkdwUDUQ2bNiAadOmYevWrVi1ahWqqqpw880348qVK0qelmQkcHCGiIgUpGiJ9xUrVtjdXrx4MRITE7Fr1y4MGTJEyVMTERGRAaiaI1JUVAQAaN68uZqn9QtyDpuwdojn8orK8c2uM6istmjdFCIiQ1Jt0TuLxYInn3wS1113HXr27Olwn4qKClRU1C2yVlxcrFbzqNYLyw7g1+MXxduMSRwTBAGPfroLKw9aa7CcvXwVT4zsrHGriIiMR7UekWnTpuHAgQP44osvnO4zd+5cxMTEiD/JyclqNc/wisurZDnOf7ecRuZF5vC4Ul1jQebFK2IQAgBvrz6GracKcCSvGFcrazRsHRGRsZgEFfrhp0+fjmXLlmHjxo1o37690/0c9YgkJyejqKgI0dHRSjfTUHKLruLfa09g0uAUdE6KQsqs5YqcJyHKjAsl1vfkjr6t8M/fpypyHqM4nl+Cm97e6HKfgSnN8dWj6Sq1iIhIf4qLixETE+PR9VvRoRlBEPDnP/8Z33//PdavX+8yCAEAs9kMs9msZJN8xrTPdmN3diG+230Wh18ZrXVz/Ia7IAQAtmddUqElRES+QdFAZNq0aViyZAmWLVuGqKgo5OXlAQBiYmIQFham5Kl93oGz1vyZq1XKDgPY1RHx86JpTOIlIpKfojki8+fPR1FREYYOHYqWLVuKP19++aWSp/UL6tX34MXXpsLJzJhTr9/S4D7OoiEi8oyigYggCA5/Jk2apORpiWRXYxGwdM/ZBvePuzYZAQEmXNuumd39T3+9FxXVTFolInKHa80YlElSOeSrHTkatsQ/LNl2GrO+2y/e7tMmBoM7xmHm6K4AgLfu62u3/w97z+G73Q0DFyIisqdaHRFSzl++3afYsZkWAVgsAp5fdlC8fUPneHwyJc1un7Zx4Xh7XB889eVe8b6C0goQEZFr7BExKK4Bo57l+3PF7f7tmmH+A/0d7ndnahuMvCZJvJ1XXK5424iIjI6BiEFV1TAQUcvyfXWByFePpCPS7Lwjce5dvdCmmXVG2I/7cjnThojIDQYi5BIvo0Dr2sAiyhyEwADXU5gToszisE1hWRXOXL6qePuIiIzMLwOR8yXl+GHvOWw4dkHrpuie9Bu9P1YRuXSlEh9tygQA3Ng1waPfaR8fgbDgQADADf9Yh3KFa70QERmZXwYih3NL8OfP9+DvPx/Ruim6d7lMnjVsjKi6xoJ+r6wSb/evN0XXlWtaRonbthL5RETUkF8GIrZv9hx2IFc2nyywuz35OtdLFEh9/IeB4vZn27JlaxMRka/xz0CkNhJhIiG5Il3R+MVbu3v1u1GhweL2BxtOytYmIiJf45eBSICfr5lCnpEmmg7qGNekY526UNrU5hAR+SS/DERsYYiFPSLkwsebswAAA1Oao1sL18tYu/OXb5QrOkdEZGR+GYhAHJrRthmG40cdSWWV1cgtshYku7lHkpu9HRvbu6W4XVpRLUu7iIh8jV8GIrZ1WhiHkDPnCuuGZWLCgl3s6dwdfVuL271axzS5TUREvsg/AxEmq5IbZwvryrPf1rdVo44x8ppEPDCoLQDg611ncLWS9USIiOrzz0Ck9v8MQ8iZhbVFzIZ1TYA5KLBRxzCZTHhyZBfx9je7z8jSNiIiX+KfgYjYJaJtO0iffjtxUay6mxQd2qRjxUeaxe2rlcwTISKqzy8DkYB6ccg/VhzBHxbvQI2FkQkBT36ZIW4/lJ7S5OM9NrQjAGBPdmGTj0VE5Gv8MhCxdYhkXryCXacv4f31J7H2yHn8duKitg1z4mIpS4SrSVqSvXNSZJOP1zc5FgDw84E8FPlxyXwiIkf8MhCRzkOdvGiHuF1ZbdGiMS59svU0rn11Nd5adUzrpvil4MCm/xOJlcy6Wb4/t8nHIyLyJX4ZiEgLq1bV6Hs45vmlBwAA76w5rnFLgBqLgK2nClBR7buzP86X1M2WGT+wrSzHDA8JEre/ZcIqEZEd/wxEJNsCM1Y9tizjHH7/n62Y/d1+rZuimDdWHBW3Hx/RSZZjSiv47jp9WZZjEhH5Cv8MRCRdInovJaLHZXG+231W6yYopvBqXQ5HpDnIxZ6eu6Zl08rDExH5Mr8MRAIkF3edxyGkspDanJCgAJPdCrpNOmZQAL59bLB4W5oMS0Tk7/wyEDHBcSSix94HHTZJl574Yg/uW7AFliZOwbbNUHprXF8ZWlVHOvvmvXUnZD22L8m6eAV7sjl8ReRP5Ol7NhiTXRyi7z4Rk8mk//EjHViWcQ4Line truncated
|
||||||
|
"text/plain": [
|
||||||
|
"<Figure size 640x480 with 1 Axes>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "display_data"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"import matplotlib.pyplot as plt\n",
|
||||||
|
"\n",
|
||||||
|
"plt.plot(total)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "eb368900-34b3-4c36-a916-62ed3a97173b",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "50253cc9-254c-4933-94bd-36c28d0f97b0",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "895a5767-2865-477f-aa21-094a711bf033",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "Python 3 (ipykernel)",
|
||||||
|
"language": "python",
|
||||||
|
"name": "python3"
|
||||||
|
},
|
||||||
|
"language_info": {
|
||||||
|
"codemirror_mode": {
|
||||||
|
"name": "ipython",
|
||||||
|
"version": 3
|
||||||
|
},
|
||||||
|
"file_extension": ".py",
|
||||||
|
"mimetype": "text/x-python",
|
||||||
|
"name": "python",
|
||||||
|
"nbconvert_exporter": "python",
|
||||||
|
"pygments_lexer": "ipython3",
|
||||||
|
"version": "3.10.10"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 5
|
||||||
|
}
|
||||||
@@ -343,7 +343,9 @@
|
|||||||
" # Increase the channel dimension from 1 to d_model\n",
|
" # Increase the channel dimension from 1 to d_model\n",
|
||||||
" u = self.fc_before(u) \n",
|
" u = self.fc_before(u) \n",
|
||||||
" # Pass through the operator\n",
|
" # Pass through the operator\n",
|
||||||
|
" #[B,1024,128] --> [B,1024,128]\n",
|
||||||
" u = self.operator(u)\n",
|
" u = self.operator(u)\n",
|
||||||
|
" \n",
|
||||||
" last_state = u[:,-1,:]\n",
|
" last_state = u[:,-1,:]\n",
|
||||||
" # Decrease the channel dimension back to 1\n",
|
" # Decrease the channel dimension back to 1\n",
|
||||||
" y = self.fc_after(last_state)\n",
|
" y = self.fc_after(last_state)\n",
|
||||||
@@ -564,10 +566,127 @@
|
|||||||
"r2_scores = fit_and_evaluate_linear_regression(outputs_and_params)"
|
"r2_scores = fit_and_evaluate_linear_regression(outputs_and_params)"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "de65d0d2-b0c6-4ac5-a87f-70fa7f90480b",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"# Autoregressive"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 120,
|
||||||
|
"id": "8ee139a6-aee5-4309-8685-bf0c28893279",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def predict_autoregressive(model, initial_data, n_steps):\n",
|
||||||
|
" model.eval()\n",
|
||||||
|
" predictions = []\n",
|
||||||
|
" current_input = initial_data\n",
|
||||||
|
" with torch.no_grad():\n",
|
||||||
|
" for _ in range(n_steps):\n",
|
||||||
|
" # Get the prediction for the next step and save it\n",
|
||||||
|
" next_output, last_state = model(current_input)\n",
|
||||||
|
" predictions.append(next_output)\n",
|
||||||
|
" #import pdb;pdb.set_trace()\n",
|
||||||
|
"\n",
|
||||||
|
" # Prepare the input for the next step\n",
|
||||||
|
" next_input = torch.cat((current_input[:, 1:, :], next_output[:, None,:]), dim=1)\n",
|
||||||
|
"\n",
|
||||||
|
" current_input = next_input\n",
|
||||||
|
"\n",
|
||||||
|
" return torch.cat(predictions, dim=1)"
|
||||||
|
]
|
||||||
|
},
|
||||||
{
|
{
|
||||||
"cell_type": "code",
|
"cell_type": "code",
|
||||||
"execution_count": null,
|
"execution_count": null,
|
||||||
"id": "61c7e1db-ffee-4ef2-a8c6-5756e326904c",
|
"id": "272040d6-4dfb-438b-a432-744e950effd1",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 167,
|
||||||
|
"id": "9f49cbf8-08e8-428a-af80-bcb2515a6327",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"torch.Size([1, 1024, 1])"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 167,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"initial_data = dataset[0]\n",
|
||||||
|
"initial_data[0][None,...].shape"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 168,
|
||||||
|
"id": "c75464e3-e44d-4325-8804-29d8081e3a45",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"autoregressive_out = predict_autoregressive(model,initial_data[0][None,...],500)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 169,
|
||||||
|
"id": "90a37c56-59f3-49bc-bfbe-d9e126c42ed1",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"total = np.concatenate([initial_data[0].squeeze().numpy() ,autoregressive_out.squeeze().numpy()])"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 170,
|
||||||
|
"id": "3d2ad0ee-7a0e-4b6f-8b40-dd565c8af84d",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"[<matplotlib.lines.Line2D at 0x7f1be29f1570>]"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 170,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAi8AAAGdCAYAAADaPpOnAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAABhFElEQVR4nO3deXhU1fkH8O8syWTfd0ggQCBA2JcA4koUrFq3alUqgorVitXqz4W2aqu1WOtuXWtxqbstbqhQZBGRyB4gLGEnISEb2bdZz++PZCYzyWxJZubemfl+nicPZO6dyXshM/POOe95j0IIIUBERETkJ5RSB0BERETUF0xeiIiIyK8weSEiIiK/wuSFiIiI/AqTFyIiIvIrTF6IiIjIrzB5ISIiIr/C5IWIiIj8ilrqADzNZDKhoqIC0dHRUCgUUodDREREbhBCoLm5GRkZGVAqnY+tBFzyUlFRgczMTKnDICIion4oKyvD4MGDnZ4TcMlLdHQ0gM6Lj4mJkTgaIiIickdTUxMyMzMt7+POBFzyYp4qiomJYfJCRETkZ9wp+WDBLhEREfkVJi9ERETkV5i8EBERkV9h8kJERER+hckLERER+RUmL0RERORXmLwQERGRX2HyQkRERH6FyQsRERH5FSYvRERE5FeYvBAREZFfYfJCREREfoXJCxEREbmlQ2/E0hV78J8dpySNI+B2lSYiIiLPq2vVYeFbW7HnVCO+LKrAnNwUxEeGShILkxciIiJyqqZZi+veKMTRmlbER4TgpesnS5a4AExeiIiIyIl2nRG3vrsdR2takREbhndvyceIlChJY2LyQkRERHYJIfDQij3YXdaAuIgQvL94BrKTIqUOiwW7REREZN/nReX4oqgCKqUCr/9qiiwSF4DJCxEREdlRVteGRz7fBwC4Z04O8oclShxRNyYvREREZEMIgYe/KEaz1oApQ+Jxx3nDpQ7JBpMXIiIisrH2QDU2lNQgRKXA338xHmqVvNIFeUVDREREkurQG/HYyv0AgFtmD8OwZGlXFtnD5IWIiIgs/l14EqV1bUiN0eCuC0ZIHY5dTF6IiIgIANCiNeDV748CAO67cBQiNfLsqMLkhYiIiAAAb/94HHWtOmQnReKqyYOkDschJi9ERESEpg493th4DABwT0GO7Ip0rck3MiIiIvKZj7aWoqnDgOHJkbh0fIbU4TjF5IWIiCjI6QwmLN90AgDw63OGQ6VUSBuQC0xeiIiIgtzXeytQ2dSBpCgNLp8k71EXgMkLERFRUBNC4I2NxwEAi84aCo1aJXFErjF5ISIiCmJbjtfhwOkmhIeoMD8/S+pw3MLkhYiIKIh9sKUUAHDFpEGIiwiVOBr3MHkhIiIKUnWtOqwqrgQAvxl1AZi8EBERBa3/7CiDzmjCuEGxyBsUK3U4bmPyQkREFISEEPhwaxkA4AY/GnUBmLwQEREFpcKjZ3C8thWRoSr8fIL8l0db82rysnHjRlx22WXIyMiAQqHA559/7vI+GzZswOTJk6HRaDBixAi8/fbb3gyRiIgoKH2yvXPU5fJJg2S7AaMjXk1eWltbMWHCBLz88stunX/8+HFccsklOP/881FUVIR77rkHt956K1avXu3NMImIiIJKq9aA1fuqAADXTBkscTR959VU6+KLL8bFF1/s9vmvvfYasrOz8cwzzwAARo8ejU2bNuG5557D3LlzvRUmERFRUPnf/kq0640YmhiBiZlxUofTZ7KqeSksLERBQYHNbXPnzkVhYaFEEREREQWez3ZVAOjs7aJQyHsfI3tkNclVWVmJ1NRUm9tSU1PR1NSE9vZ2hIeH97qPVquFVqu1fN/U1OT1OImIiPxVdXMHNh2uAQBcMXGQxNH0j6xGXvpj2bJliI2NtXxlZmZKHRIREZFsfbX7NEwCmJgZh6FJkVKH0y+ySl7S0tJQVVVlc1tVVRViYmLsjroAwNKlS9HY2Gj5Kisr80WoREREfunzXeUAgCsn+eeoCyCzaaOZM2fim2++sbltzZo1mDlzpsP7aDQaaDQab4dGRETk90rPtGFveSOUCuCS8elSh9NvXh15aWlpQVFREYqKigB0LoUuKipCaWnnJlBLly7FggULLOfffvvtOHbsGB544AEcPHgQr7zyCj755BP87ne/82aYREREQeHb4tMAgBnDEpEU5b8f/L2avGzfvh2TJk3CpEmTAAD33nsvJk2ahEceeQQAcPr0aUsiAwDZ2dn4+uuvsWbNGkyYMAHPPPMM3nzzTS6TJiIi8oBvuzZhvDgvTeJIBkYhhBBSB+FJTU1NiI2NRWNjI2JiYqQOh4iISBYqGtox68l1UCiALUvnICUmTOqQbPTl/VtWBbtERETkHau6Rl2mDomXXeLSV0xeiIiIgoA5eZmX57+FumZMXoiIiAJcdXMHtp2sAwDM8/N6F4DJCxERUcBbva8KQgATMuMwKM5+3zR/wuSFiIgowK3Z39kAdt5Y/x91AZi8EBERBbRWrQE/HT0DALhwTIrE0XgGkxciIqIAtulILXRGE7ISIjA8OUrqcDyCyQsREVEAW3egGgAwZ3QKFAqFxNF4BpMXIiKiAGUyCawr6UpeclMljsZzmLwQEREFqOKKRtQ0axEZqsL07ASpw/EYJi9EREQBam3XlNE5I5MRqg6ct/zAuRIiIiKyse5gZ/JyQW5grDIyY/JCREQUgKqaOrC3vBEKBXA+kxciIhqI7SfqcMHTG7DxUA0AYENJNT7dXiZxVBRovu/6/Ro/OA5JURqJo/EsJi9ERD52/T9/wrHaVixYvhVagxEL39qG+/+zB6cb26UOjQLID4drAQDnjkyWOBLPY/JCRORjeqOw/P3NH45b/q7VmyCEsHcXoj4xmQQ2He4ceTknJ0niaDyPyQsRkYRqW7SWvz/wnz045+/r0ao1SBgRBYJ9FU2ob9MjWqPGhMw4qcPxOCYvREQ+YjL1HlUxWt229UQdyurasXJPhS/DogC0sWvUZebwRISoAu+tPvCuiPrNaBJo09l+4uMQNpFnPPu/Ekz5yxqU1bXZ3G6wk9CYb2rXGX0RGgWgH7qSl7MDsN4FYPJCVq5+dTPGPLIaZ7qGsc+0aHH+0xtw14e7ep37v32V+HfhCR9HSOS/Xlx3BPVtejy35pDN7R9sKbV7/svrj2D0I6uw9kCVL8KjANKqNWDHyXoAgVnvAjB5IStFZQ0AupsaPbvmEE6cacNXu3sPYd/27x14+It9luyeiLq164w4Udtq/6Ab++IJAfx9dQkA4JZ3tuOCpzdgy7EzHoyQAtlPx85AbxTISojAkMRIqcPxCiYv1It5EPtwVYvLc2/819Zew+DW2nVGmzl9omBw6Us/4LynN1g+/VpT9mNX32O1rfjlGz/h1ne2Wz5k1LfqcOs727Cq+PRAw6UAY14ifXaAjroATF7Inq5co9RBUqI3mmy+7zkMbtbYpsfoR1bh6lc3ezQ8Irk7WtM56mJv1LLvqUu37w5U4YqXfwQAPLOmBN8dqMbt7+0cwCNSIDIX656dE5j1LgCTFwKgM5jw8bbueXfRlb3Utensnt+uty0iXLGr3GZpZ1ldGx79ohjvbTkJoHs6iigYlFQ2W/5uMJl6He/HwItdZ1rsPz8puJU3tONYTStUSgVmDk+UOhyvUUsdAEnvte+P4lmr0RPzAiOdofcLLwB06HuvgPjPjlO4adZQAMBNy7fimKP5fqIAt3TFHsvf7U2ZKtwYexFwPtWqM5j6Nf1EgW/zkc4po/GDYxEbHiJxNN7DkRfC+pJqm+8FgMrGDofnd+h6JzXWiQ4TFwpm1kufDcbeScjHbuxhZK8fjLXb39sBpZLJC/VW2FXYPXNY4I66AExegp7WYMSu0gab24QANh+ttbmtvKEdf/hsL45UN6PD0HvkpbKpo9d9iIKRdWuk/har6+wkPdbWHayGirkL9SCEwE9Hu5KXAJ4yAjhtFPS+2NW7oFBAoK7Vdj59wb+24GhNK9530JPiX5uO41+bjmPlXbPtHjeZBD8pUtAx9rPJozvbAzS06/v12BS4SuvaUNHYgRCVAlOGxEsdjldx5CXI2RtFEQKo71Gsa1494cpBq2JFa3o7hYtEgWhveaPl7/srmvr1GM86WMFnbUMJeyyRrZ+6powmDI5DRGhgj00weQlySVGaXrcJAPVt/ftUV91sv1bG3tw/UaA7XO26VxKRpxQGyZQRwOQlaOgMJqw/WN1rONpkb1hbCNS39m8Z5lOrSuzezuSFggEbMpJUhBBBU6wLMHkJGk+tOohFb2/DHe/vtKxkeGntYSz5oPe+RZ0jL57tIaEzctqIAp+9vi7c3JR84cSZNlQ1aRGqUmJygNe7AExegsYHWzsLbTceqsGkx9egpLIZzziYVxcCqG/1bDGgvRd1okBjb4SxXW/s1ZXa8z+Xz69gZ54ympgVh7AQlcTReB+TlyChtlrp09iux8NfFDs896V1h1FSZb/wtr8e+M8e1ycR+blDdp43b/5wHPd9sturP3fFznIAwLuFJ3DvJ0WcvgpCwTRlBDB5CRohKtv/ameLlmvdbDsepXG/mt28URhRoPp6z2lc+UrvfbyeXXMIX9rZ48iT9pQ3AAAe+WIfVuwsx9oDVV79eSQvQoigKtYFmLwEDVWPHiue6Cx+QW7KwB+EKADUtepw5wfSbZCYEBFq832LG31iKHAcrWlFbYsWGrUSEzPjpA7HJwJ7IThZqHskL+UN7QN+zL4OTNe2aLFydwX0RoHF5wwb8M8nkpoQAi1aA2qatR55vFCVsl/F7Wmx4R75+eSfzP1dJmfFB0W9C8CRl6Ch7jFtVFbngeSlj6soLntpE/701X488c0B7D3V6PoORDLzxsajuGn5VmgNRrRqDbjzg50Y96f/9dofzJWFXZuY9tTf0UylArj1nW2W780jq0IINHWwE2+g23aiDgCQPyxB4kh8h8lLkFB7YSOUvq4APW212eNl/9iEWcvLine truncated
|
||||||
|
"text/plain": [
|
||||||
|
"<Figure size 640x480 with 1 Axes>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "display_data"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"import matplotlib.pyplot as plt\n",
|
||||||
|
"\n",
|
||||||
|
"plt.plot(total)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "eb368900-34b3-4c36-a916-62ed3a97173b",
|
||||||
"metadata": {},
|
"metadata": {},
|
||||||
"outputs": [],
|
"outputs": [],
|
||||||
"source": []
|
"source": []
|
||||||
|
|||||||
@@ -0,0 +1,850 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 32,
|
||||||
|
"id": "6c6e33cb-72f9-42fa-936a-33b5fe338d15",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"torch.Size([1, 1024, 128]) torch.Size([1, 1024, 128])\n",
|
||||||
|
"Causality check: gradients should not flow \"from future to past\"\n",
|
||||||
|
"tensor(-6.7998e-10) tensor(0.1017)\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"# %load standalone_hyena.py\n",
|
||||||
|
"\"\"\"\n",
|
||||||
|
"Simplified standalone version of Hyena: https://arxiv.org/abs/2302.10866, designed for quick experimentation.\n",
|
||||||
|
"A complete version is available under `src.models.sequence.hyena`.\n",
|
||||||
|
"\"\"\"\n",
|
||||||
|
"\n",
|
||||||
|
"import math\n",
|
||||||
|
"import torch\n",
|
||||||
|
"import torch.nn as nn\n",
|
||||||
|
"import torch.nn.functional as F\n",
|
||||||
|
"from einops import rearrange\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"def fftconv(u, k, D):\n",
|
||||||
|
" seqlen = u.shape[-1]\n",
|
||||||
|
" fft_size = 2 * seqlen\n",
|
||||||
|
" \n",
|
||||||
|
" k_f = torch.fft.rfft(k, n=fft_size) / fft_size\n",
|
||||||
|
" u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)\n",
|
||||||
|
" \n",
|
||||||
|
" if len(u.shape) > 3: k_f = k_f.unsqueeze(1)\n",
|
||||||
|
" y = torch.fft.irfft(u_f * k_f, n=fft_size, norm='forward')[..., :seqlen]\n",
|
||||||
|
"\n",
|
||||||
|
" out = y + u * D.unsqueeze(-1)\n",
|
||||||
|
" return out.to(dtype=u.dtype)\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"@torch.jit.script \n",
|
||||||
|
"def mul_sum(q, y):\n",
|
||||||
|
" return (q * y).sum(dim=1)\n",
|
||||||
|
"\n",
|
||||||
|
"class OptimModule(nn.Module):\n",
|
||||||
|
" \"\"\" Interface for Module that allows registering buffers/parameters with configurable optimizer hyperparameters \"\"\"\n",
|
||||||
|
"\n",
|
||||||
|
" def register(self, name, tensor, lr=None, wd=0.0):\n",
|
||||||
|
" \"\"\"Register a tensor with a configurable learning rate and 0 weight decay\"\"\"\n",
|
||||||
|
"\n",
|
||||||
|
" if lr == 0.0:\n",
|
||||||
|
" self.register_buffer(name, tensor)\n",
|
||||||
|
" else:\n",
|
||||||
|
" self.register_parameter(name, nn.Parameter(tensor))\n",
|
||||||
|
"\n",
|
||||||
|
" optim = {}\n",
|
||||||
|
" if lr is not None: optim[\"lr\"] = lr\n",
|
||||||
|
" if wd is not None: optim[\"weight_decay\"] = wd\n",
|
||||||
|
" setattr(getattr(self, name), \"_optim\", optim)\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
"class Sin(nn.Module):\n",
|
||||||
|
" def __init__(self, dim, w=10, train_freq=True):\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" self.freq = nn.Parameter(w * torch.ones(1, dim)) if train_freq else w * torch.ones(1, dim)\n",
|
||||||
|
"\n",
|
||||||
|
" def forward(self, x):\n",
|
||||||
|
" return torch.sin(self.freq * x)\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
"class PositionalEmbedding(OptimModule):\n",
|
||||||
|
" def __init__(self, emb_dim: int, seq_len: int, lr_pos_emb: float=1e-5, **kwargs): \n",
|
||||||
|
" \"\"\"Complex exponential positional embeddings for Hyena filters.\"\"\" \n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" \n",
|
||||||
|
" self.seq_len = seq_len\n",
|
||||||
|
" # The time embedding fed to the filteres is normalized so that t_f = 1\n",
|
||||||
|
" t = torch.linspace(0, 1, self.seq_len)[None, :, None] # 1, L, 1\n",
|
||||||
|
" \n",
|
||||||
|
" if emb_dim > 1:\n",
|
||||||
|
" bands = (emb_dim - 1) // 2 \n",
|
||||||
|
" # To compute the right embeddings we use the \"proper\" linspace \n",
|
||||||
|
" t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]\n",
|
||||||
|
" w = 2 * math.pi * t_rescaled / seq_len # 1, L, 1 \n",
|
||||||
|
" \n",
|
||||||
|
" f = torch.linspace(1e-4, bands - 1, bands)[None, None] \n",
|
||||||
|
" z = torch.exp(-1j * f * w)\n",
|
||||||
|
" z = torch.cat([t, z.real, z.imag], dim=-1)\n",
|
||||||
|
" self.register(\"z\", z, lr=lr_pos_emb) \n",
|
||||||
|
" self.register(\"t\", t, lr=0.0)\n",
|
||||||
|
" \n",
|
||||||
|
" def forward(self, L):\n",
|
||||||
|
" return self.z[:, :L], self.t[:, :L]\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
"class ExponentialModulation(OptimModule):\n",
|
||||||
|
" def __init__(\n",
|
||||||
|
" self,\n",
|
||||||
|
" d_model,\n",
|
||||||
|
" fast_decay_pct=0.3,\n",
|
||||||
|
" slow_decay_pct=1.5,\n",
|
||||||
|
" target=1e-2,\n",
|
||||||
|
" modulation_lr=0.0,\n",
|
||||||
|
" modulate: bool=True,\n",
|
||||||
|
" shift: float = 0.0,\n",
|
||||||
|
" **kwargs\n",
|
||||||
|
" ):\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" self.modulate = modulate\n",
|
||||||
|
" self.shift = shift\n",
|
||||||
|
" max_decay = math.log(target) / fast_decay_pct\n",
|
||||||
|
" min_decay = math.log(target) / slow_decay_pct\n",
|
||||||
|
" deltas = torch.linspace(min_decay, max_decay, d_model)[None, None]\n",
|
||||||
|
" self.register(\"deltas\", deltas, lr=modulation_lr)\n",
|
||||||
|
" \n",
|
||||||
|
" def forward(self, t, x):\n",
|
||||||
|
" if self.modulate:\n",
|
||||||
|
" decay = torch.exp(-t * self.deltas.abs()) \n",
|
||||||
|
" x = x * (decay + self.shift)\n",
|
||||||
|
" return x \n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"class HyenaFilter(OptimModule):\n",
|
||||||
|
" def __init__(\n",
|
||||||
|
" self, \n",
|
||||||
|
" d_model,\n",
|
||||||
|
" emb_dim=3, # dim of input to MLP, augments with positional encoding\n",
|
||||||
|
" order=16, # width of the implicit MLP \n",
|
||||||
|
" fused_fft_conv=False,\n",
|
||||||
|
" seq_len=1024, \n",
|
||||||
|
" lr=1e-3, \n",
|
||||||
|
" lr_pos_emb=1e-5,\n",
|
||||||
|
" dropout=0.0, \n",
|
||||||
|
" w=1, # frequency of periodic activations \n",
|
||||||
|
" wd=0, # weight decay of kernel parameters \n",
|
||||||
|
" bias=True,\n",
|
||||||
|
" num_inner_mlps=2,\n",
|
||||||
|
" normalized=False,\n",
|
||||||
|
" **kwargs\n",
|
||||||
|
" ):\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" Implicit long filter with modulation.\n",
|
||||||
|
" \n",
|
||||||
|
" Args:\n",
|
||||||
|
" d_model: number of channels in the input\n",
|
||||||
|
" emb_dim: dimension of the positional encoding (`emb_dim` - 1) // 2 is the number of bands\n",
|
||||||
|
" order: width of the FFN\n",
|
||||||
|
" num_inner_mlps: number of inner linear layers inside filter MLP\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" self.d_model = d_model\n",
|
||||||
|
" self.use_bias = bias\n",
|
||||||
|
" self.fused_fft_conv = fused_fft_conv\n",
|
||||||
|
" self.bias = nn.Parameter(torch.randn(self.d_model))\n",
|
||||||
|
" self.dropout = nn.Dropout(dropout)\n",
|
||||||
|
" \n",
|
||||||
|
" act = Sin(dim=order, w=w)\n",
|
||||||
|
" self.emb_dim = emb_dim\n",
|
||||||
|
" assert emb_dim % 2 != 0 and emb_dim >= 3, \"emb_dim must be odd and greater or equal to 3 (time, sine and cosine)\"\n",
|
||||||
|
" self.seq_len = seq_len\n",
|
||||||
|
" \n",
|
||||||
|
" self.pos_emb = PositionalEmbedding(emb_dim, seq_len, lr_pos_emb)\n",
|
||||||
|
" \n",
|
||||||
|
" self.implicit_filter = nn.Sequential(\n",
|
||||||
|
" nn.Linear(emb_dim, order),\n",
|
||||||
|
" act,\n",
|
||||||
|
" )\n",
|
||||||
|
" for i in range(num_inner_mlps):\n",
|
||||||
|
" self.implicit_filter.append(nn.Linear(order, order))\n",
|
||||||
|
" self.implicit_filter.append(act)\n",
|
||||||
|
"\n",
|
||||||
|
" self.implicit_filter.append(nn.Linear(order, d_model, bias=False))\n",
|
||||||
|
" \n",
|
||||||
|
" self.modulation = ExponentialModulation(d_model, **kwargs)\n",
|
||||||
|
" \n",
|
||||||
|
" self.normalized = normalized\n",
|
||||||
|
" for c in self.implicit_filter.children():\n",
|
||||||
|
" for name, v in c.state_dict().items(): \n",
|
||||||
|
" optim = {\"weight_decay\": wd, \"lr\": lr}\n",
|
||||||
|
" setattr(getattr(c, name), \"_optim\", optim)\n",
|
||||||
|
"\n",
|
||||||
|
" def filter(self, L, *args, **kwargs):\n",
|
||||||
|
" z, t = self.pos_emb(L)\n",
|
||||||
|
" h = self.implicit_filter(z)\n",
|
||||||
|
" h = self.modulation(t, h)\n",
|
||||||
|
" return h\n",
|
||||||
|
"\n",
|
||||||
|
" def forward(self, x, L, k=None, bias=None, *args, **kwargs):\n",
|
||||||
|
" if k is None: k = self.filter(L)\n",
|
||||||
|
" \n",
|
||||||
|
" # Ensure compatibility with filters that return a tuple \n",
|
||||||
|
" k = k[0] if type(k) is tuple else k \n",
|
||||||
|
"\n",
|
||||||
|
" y = fftconv(x, k, bias)\n",
|
||||||
|
" return y\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
"class HyenaOperator(nn.Module):\n",
|
||||||
|
" def __init__(\n",
|
||||||
|
" self,\n",
|
||||||
|
" d_model,\n",
|
||||||
|
" l_max,\n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64,\n",
|
||||||
|
" dropout=0.0, \n",
|
||||||
|
" filter_dropout=0.0, \n",
|
||||||
|
" **filter_args,\n",
|
||||||
|
" ):\n",
|
||||||
|
" r\"\"\"\n",
|
||||||
|
" Hyena operator described in the paper https://arxiv.org/pdf/2302.10866.pdf\n",
|
||||||
|
" \n",
|
||||||
|
" Args:\n",
|
||||||
|
" d_model (int): Dimension of the input and output embeddings (width of the layer)\n",
|
||||||
|
" l_max: (int): Maximum input sequence length. Defaults to None\n",
|
||||||
|
" order: (int): Depth of the Hyena recurrence. Defaults to 2\n",
|
||||||
|
" dropout: (float): Dropout probability. Defaults to 0.0\n",
|
||||||
|
" filter_dropout: (float): Dropout probability for the filter. Defaults to 0.0\n",
|
||||||
|
" \"\"\"\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
" self.d_model = d_model\n",
|
||||||
|
" self.l_max = l_max\n",
|
||||||
|
" self.order = order\n",
|
||||||
|
" inner_width = d_model * (order + 1)\n",
|
||||||
|
" self.dropout = nn.Dropout(dropout)\n",
|
||||||
|
" self.in_proj = nn.Linear(d_model, inner_width)\n",
|
||||||
|
" self.out_proj = nn.Linear(d_model, d_model)\n",
|
||||||
|
" \n",
|
||||||
|
" self.short_filter = nn.Conv1d(\n",
|
||||||
|
" inner_width, \n",
|
||||||
|
" inner_width, \n",
|
||||||
|
" 3,\n",
|
||||||
|
" padding=2,\n",
|
||||||
|
" groups=inner_width\n",
|
||||||
|
" )\n",
|
||||||
|
" self.filter_fn = HyenaFilter(\n",
|
||||||
|
" d_model * (order - 1), \n",
|
||||||
|
" order=filter_order, \n",
|
||||||
|
" seq_len=l_max,\n",
|
||||||
|
" channels=1, \n",
|
||||||
|
" dropout=filter_dropout, \n",
|
||||||
|
" **filter_args\n",
|
||||||
|
" ) \n",
|
||||||
|
"\n",
|
||||||
|
" def forward(self, u, *args, **kwargs):\n",
|
||||||
|
" l = u.size(-2)\n",
|
||||||
|
" l_filter = min(l, self.l_max)\n",
|
||||||
|
" u = self.in_proj(u)\n",
|
||||||
|
" u = rearrange(u, 'b l d -> b d l')\n",
|
||||||
|
" \n",
|
||||||
|
" uc = self.short_filter(u)[...,:l_filter] \n",
|
||||||
|
" *x, v = uc.split(self.d_model, dim=1)\n",
|
||||||
|
" \n",
|
||||||
|
" k = self.filter_fn.filter(l_filter)[0]\n",
|
||||||
|
" k = rearrange(k, 'l (o d) -> o d l', o=self.order - 1)\n",
|
||||||
|
" bias = rearrange(self.filter_fn.bias, '(o d) -> o d', o=self.order - 1)\n",
|
||||||
|
" \n",
|
||||||
|
" for o, x_i in enumerate(reversed(x[1:])):\n",
|
||||||
|
" v = self.dropout(v * x_i)\n",
|
||||||
|
" v = self.filter_fn(v, l_filter, k=k[o], bias=bias[o])\n",
|
||||||
|
"\n",
|
||||||
|
" y = rearrange(v * x[0], 'b d l -> b l d')\n",
|
||||||
|
"\n",
|
||||||
|
" y = self.out_proj(y)\n",
|
||||||
|
" return y\n",
|
||||||
|
"\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
"if __name__ == \"__main__\":\n",
|
||||||
|
" layer = HyenaOperator(\n",
|
||||||
|
" \n",
|
||||||
|
" d_model=128, \n",
|
||||||
|
" l_max=1024, \n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64\n",
|
||||||
|
" )\n",
|
||||||
|
" x = torch.randn(1, 1024, 128, requires_grad=True)\n",
|
||||||
|
" y = layer(x)\n",
|
||||||
|
" \n",
|
||||||
|
" print(x.shape, y.shape)\n",
|
||||||
|
" \n",
|
||||||
|
" grad = torch.autograd.grad(y[:, 10, :].sum(), x)[0]\n",
|
||||||
|
" print('Causality check: gradients should not flow \"from future to past\"')\n",
|
||||||
|
" print(grad[0, 11, :].sum(), grad[0, 9, :].sum())\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 368,
|
||||||
|
"id": "032ef08a-8cc6-491a-9eb8-4a6b3f2d165e",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"torch.Size([1, 1023, 1]) torch.Size([1, 500])\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"class HyenaOperatorAutoregressive1D(nn.Module):\n",
|
||||||
|
" def __init__(\n",
|
||||||
|
" self,\n",
|
||||||
|
" d_model,\n",
|
||||||
|
" l_max,\n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64,\n",
|
||||||
|
" dropout=0.0, \n",
|
||||||
|
" filter_dropout=0.0, \n",
|
||||||
|
" **filter_args,\n",
|
||||||
|
" ):\n",
|
||||||
|
" super().__init__()\n",
|
||||||
|
"\n",
|
||||||
|
" self.l_max = l_max\n",
|
||||||
|
" self.d_model = d_model\n",
|
||||||
|
" self.l_max = l_max\n",
|
||||||
|
" self.order = order\n",
|
||||||
|
" inner_width = d_model * (order + 1)\n",
|
||||||
|
"\n",
|
||||||
|
" self.dropout = nn.Dropout(dropout)\n",
|
||||||
|
" self.in_proj = nn.Linear(d_model, inner_width)\n",
|
||||||
|
" self.out_proj = nn.Linear(d_model, d_model)\n",
|
||||||
|
" self.fc_before = nn.Linear(1, d_model) # Fully connected layer before the main layer\n",
|
||||||
|
" self.fc_after = nn.Linear(d_model, 500) # Fully connected layer after the main layer\n",
|
||||||
|
"\n",
|
||||||
|
" self.operator = HyenaOperator(\n",
|
||||||
|
" d_model=d_model,\n",
|
||||||
|
" l_max=l_max,\n",
|
||||||
|
" order=order, \n",
|
||||||
|
" filter_order=filter_order,\n",
|
||||||
|
" dropout=dropout, \n",
|
||||||
|
" filter_dropout=filter_dropout, \n",
|
||||||
|
" **filter_args,\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" def forward(self, u, *args, **kwargs):\n",
|
||||||
|
" # Increase the channel dimension from 1 to d_model\n",
|
||||||
|
" u = self.fc_before(u) \n",
|
||||||
|
" # Pass through the operator\n",
|
||||||
|
" #[B,1024,128] --> [B,1024,128]\n",
|
||||||
|
" u = self.operator(u)\n",
|
||||||
|
" \n",
|
||||||
|
" last_state = u[:,-1,:]\n",
|
||||||
|
" # Decrease the channel dimension back to 1\n",
|
||||||
|
" y = self.fc_after(last_state)\n",
|
||||||
|
" return y,last_state\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"if __name__ == \"__main__\":\n",
|
||||||
|
" layer = HyenaOperatorAutoregressive1D(\n",
|
||||||
|
" d_model=128, \n",
|
||||||
|
" l_max=1024, \n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64\n",
|
||||||
|
" )\n",
|
||||||
|
"\n",
|
||||||
|
" x = torch.randn(1, 1023, 1, requires_grad=True) # 1D time series input\n",
|
||||||
|
" y, last_state = layer(x)\n",
|
||||||
|
"\n",
|
||||||
|
" #import pdb;pdb.set_trace()\n",
|
||||||
|
" print(x.shape, y.shape) # should now be [1, 1024, 1]\n",
|
||||||
|
"\n",
|
||||||
|
" #grad = torch.autograd.grad(y[:, 10, 0].sum(), x)[0]\n",
|
||||||
|
" #print('Causality check: gradients should not flow \"from future to past\"')\n",
|
||||||
|
" #print(grad[0, 11, 0].sum(), grad[0, 9, 0].sum())"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 369,
|
||||||
|
"id": "80cde67b-992f-4cb0-8824-4a6b7e4984ca",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"Train Epoch: 1 [0/640 (0%)]\tLoss: 0.569402\n",
|
||||||
|
"Train Epoch: 2 [0/640 (0%)]\tLoss: 0.459675\n",
|
||||||
|
"Train Epoch: 3 [0/640 (0%)]\tLoss: 0.455738\n",
|
||||||
|
"Train Epoch: 4 [0/640 (0%)]\tLoss: 0.334769\n",
|
||||||
|
"Train Epoch: 5 [0/640 (0%)]\tLoss: 0.400367\n",
|
||||||
|
"Train Epoch: 6 [0/640 (0%)]\tLoss: 0.353075\n",
|
||||||
|
"Train Epoch: 7 [0/640 (0%)]\tLoss: 0.372000\n",
|
||||||
|
"Train Epoch: 8 [0/640 (0%)]\tLoss: 0.354896\n",
|
||||||
|
"Train Epoch: 9 [0/640 (0%)]\tLoss: 0.370209\n",
|
||||||
|
"Train Epoch: 10 [0/640 (0%)]\tLoss: 0.323172\n",
|
||||||
|
"Train Epoch: 11 [0/640 (0%)]\tLoss: 0.324925\n",
|
||||||
|
"Train Epoch: 12 [0/640 (0%)]\tLoss: 0.319851\n",
|
||||||
|
"Train Epoch: 13 [0/640 (0%)]\tLoss: 0.267247\n",
|
||||||
|
"Train Epoch: 14 [0/640 (0%)]\tLoss: 0.307432\n",
|
||||||
|
"Train Epoch: 15 [0/640 (0%)]\tLoss: 0.270438\n",
|
||||||
|
"Train Epoch: 16 [0/640 (0%)]\tLoss: 0.274847\n",
|
||||||
|
"Train Epoch: 17 [0/640 (0%)]\tLoss: 0.283570\n",
|
||||||
|
"Train Epoch: 18 [0/640 (0%)]\tLoss: 0.282435\n",
|
||||||
|
"Train Epoch: 19 [0/640 (0%)]\tLoss: 0.325277\n",
|
||||||
|
"Train Epoch: 20 [0/640 (0%)]\tLoss: 0.326561\n",
|
||||||
|
"Train Epoch: 21 [0/640 (0%)]\tLoss: 0.305762\n",
|
||||||
|
"Train Epoch: 22 [0/640 (0%)]\tLoss: 0.298317\n",
|
||||||
|
"Train Epoch: 23 [0/640 (0%)]\tLoss: 0.277244\n",
|
||||||
|
"Train Epoch: 24 [0/640 (0%)]\tLoss: 0.271200\n",
|
||||||
|
"Train Epoch: 25 [0/640 (0%)]\tLoss: 0.246922\n",
|
||||||
|
"Train Epoch: 26 [0/640 (0%)]\tLoss: 0.231166\n",
|
||||||
|
"Train Epoch: 27 [0/640 (0%)]\tLoss: 0.301950\n",
|
||||||
|
"Train Epoch: 28 [0/640 (0%)]\tLoss: 0.272427\n",
|
||||||
|
"Train Epoch: 29 [0/640 (0%)]\tLoss: 0.278342\n",
|
||||||
|
"Train Epoch: 30 [0/640 (0%)]\tLoss: 0.246955\n",
|
||||||
|
"Train Epoch: 31 [0/640 (0%)]\tLoss: 0.304109\n",
|
||||||
|
"Train Epoch: 32 [0/640 (0%)]\tLoss: 0.254574\n",
|
||||||
|
"Train Epoch: 33 [0/640 (0%)]\tLoss: 0.234942\n",
|
||||||
|
"Train Epoch: 34 [0/640 (0%)]\tLoss: 0.269669\n",
|
||||||
|
"Train Epoch: 35 [0/640 (0%)]\tLoss: 0.239477\n",
|
||||||
|
"Train Epoch: 36 [0/640 (0%)]\tLoss: 0.282838\n",
|
||||||
|
"Train Epoch: 37 [0/640 (0%)]\tLoss: 0.264821\n",
|
||||||
|
"Train Epoch: 38 [0/640 (0%)]\tLoss: 0.236031\n",
|
||||||
|
"Train Epoch: 39 [0/640 (0%)]\tLoss: 0.239743\n",
|
||||||
|
"Train Epoch: 40 [0/640 (0%)]\tLoss: 0.289715\n",
|
||||||
|
"Train Epoch: 41 [0/640 (0%)]\tLoss: 0.240897\n",
|
||||||
|
"Train Epoch: 42 [0/640 (0%)]\tLoss: 0.212328\n",
|
||||||
|
"Train Epoch: 43 [0/640 (0%)]\tLoss: 0.253719\n",
|
||||||
|
"Train Epoch: 44 [0/640 (0%)]\tLoss: 0.210382\n",
|
||||||
|
"Train Epoch: 45 [0/640 (0%)]\tLoss: 0.284465\n",
|
||||||
|
"Train Epoch: 46 [0/640 (0%)]\tLoss: 0.292590\n",
|
||||||
|
"Train Epoch: 47 [0/640 (0%)]\tLoss: 0.200158\n",
|
||||||
|
"Train Epoch: 48 [0/640 (0%)]\tLoss: 0.269204\n",
|
||||||
|
"Train Epoch: 49 [0/640 (0%)]\tLoss: 0.213730\n",
|
||||||
|
"Train Epoch: 50 [0/640 (0%)]\tLoss: 0.268548\n",
|
||||||
|
"Train Epoch: 51 [0/640 (0%)]\tLoss: 0.186853\n",
|
||||||
|
"Train Epoch: 52 [0/640 (0%)]\tLoss: 0.263178\n",
|
||||||
|
"Train Epoch: 53 [0/640 (0%)]\tLoss: 0.232999\n",
|
||||||
|
"Train Epoch: 54 [0/640 (0%)]\tLoss: 0.265539\n",
|
||||||
|
"Train Epoch: 55 [0/640 (0%)]\tLoss: 0.227278\n",
|
||||||
|
"Train Epoch: 56 [0/640 (0%)]\tLoss: 0.266715\n",
|
||||||
|
"Train Epoch: 57 [0/640 (0%)]\tLoss: 0.243662\n",
|
||||||
|
"Train Epoch: 58 [0/640 (0%)]\tLoss: 0.231188\n",
|
||||||
|
"Train Epoch: 59 [0/640 (0%)]\tLoss: 0.234178\n",
|
||||||
|
"Train Epoch: 60 [0/640 (0%)]\tLoss: 0.270513\n",
|
||||||
|
"Train Epoch: 61 [0/640 (0%)]\tLoss: 0.240779\n",
|
||||||
|
"Train Epoch: 62 [0/640 (0%)]\tLoss: 0.278103\n",
|
||||||
|
"Train Epoch: 63 [0/640 (0%)]\tLoss: 0.235614\n",
|
||||||
|
"Train Epoch: 64 [0/640 (0%)]\tLoss: 0.248226\n",
|
||||||
|
"Train Epoch: 65 [0/640 (0%)]\tLoss: 0.242527\n",
|
||||||
|
"Train Epoch: 66 [0/640 (0%)]\tLoss: 0.312065\n",
|
||||||
|
"Train Epoch: 67 [0/640 (0%)]\tLoss: 0.201206\n",
|
||||||
|
"Train Epoch: 68 [0/640 (0%)]\tLoss: 0.263827\n",
|
||||||
|
"Train Epoch: 69 [0/640 (0%)]\tLoss: 0.286222\n",
|
||||||
|
"Train Epoch: 70 [0/640 (0%)]\tLoss: 0.279344\n",
|
||||||
|
"Train Epoch: 71 [0/640 (0%)]\tLoss: 0.263106\n",
|
||||||
|
"Train Epoch: 72 [0/640 (0%)]\tLoss: 0.247235\n",
|
||||||
|
"Train Epoch: 73 [0/640 (0%)]\tLoss: 0.243001\n",
|
||||||
|
"Train Epoch: 74 [0/640 (0%)]\tLoss: 0.231101\n",
|
||||||
|
"Train Epoch: 75 [0/640 (0%)]\tLoss: 0.290737\n",
|
||||||
|
"Train Epoch: 76 [0/640 (0%)]\tLoss: 0.274286\n",
|
||||||
|
"Train Epoch: 77 [0/640 (0%)]\tLoss: 0.254139\n",
|
||||||
|
"Train Epoch: 78 [0/640 (0%)]\tLoss: 0.296462\n",
|
||||||
|
"Train Epoch: 79 [0/640 (0%)]\tLoss: 0.223839\n",
|
||||||
|
"Train Epoch: 80 [0/640 (0%)]\tLoss: 0.282157\n",
|
||||||
|
"Train Epoch: 81 [0/640 (0%)]\tLoss: 0.193232\n",
|
||||||
|
"Train Epoch: 82 [0/640 (0%)]\tLoss: 0.319076\n",
|
||||||
|
"Train Epoch: 83 [0/640 (0%)]\tLoss: 0.195717\n",
|
||||||
|
"Train Epoch: 84 [0/640 (0%)]\tLoss: 0.243448\n",
|
||||||
|
"Train Epoch: 85 [0/640 (0%)]\tLoss: 0.218913\n",
|
||||||
|
"Train Epoch: 86 [0/640 (0%)]\tLoss: 0.267013\n",
|
||||||
|
"Train Epoch: 87 [0/640 (0%)]\tLoss: 0.243178\n",
|
||||||
|
"Train Epoch: 88 [0/640 (0%)]\tLoss: 0.211878\n",
|
||||||
|
"Train Epoch: 89 [0/640 (0%)]\tLoss: 0.294469\n",
|
||||||
|
"Train Epoch: 90 [0/640 (0%)]\tLoss: 0.270843\n",
|
||||||
|
"Train Epoch: 91 [0/640 (0%)]\tLoss: 0.212928\n",
|
||||||
|
"Train Epoch: 92 [0/640 (0%)]\tLoss: 0.284125\n",
|
||||||
|
"Train Epoch: 93 [0/640 (0%)]\tLoss: 0.236006\n",
|
||||||
|
"Train Epoch: 94 [0/640 (0%)]\tLoss: 0.260203\n",
|
||||||
|
"Train Epoch: 95 [0/640 (0%)]\tLoss: 0.270325\n",
|
||||||
|
"Train Epoch: 96 [0/640 (0%)]\tLoss: 0.199420\n",
|
||||||
|
"Train Epoch: 97 [0/640 (0%)]\tLoss: 0.313867\n",
|
||||||
|
"Train Epoch: 98 [0/640 (0%)]\tLoss: 0.285310\n",
|
||||||
|
"Train Epoch: 99 [0/640 (0%)]\tLoss: 0.284724\n",
|
||||||
|
"Train Epoch: 100 [0/640 (0%)]\tLoss: 0.256413\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"import torch\n",
|
||||||
|
"import torch.optim as optim\n",
|
||||||
|
"import torch.nn.functional as F\n",
|
||||||
|
"from torch.utils.data import DataLoader, Dataset\n",
|
||||||
|
"import numpy as np\n",
|
||||||
|
"\n",
|
||||||
|
"def generate_sine_with_noise(n_points, frequency, phase, amplitude, noise_sd):\n",
|
||||||
|
" # Generate an array of points from 0 to 2*pi\n",
|
||||||
|
" x = np.linspace(0, 2*np.pi, n_points)\n",
|
||||||
|
" \n",
|
||||||
|
" # Generate the sine wave\n",
|
||||||
|
" sine_wave = amplitude * np.sin(frequency * x + phase)\n",
|
||||||
|
" \n",
|
||||||
|
" # Generate Gaussian noise\n",
|
||||||
|
" noise = np.random.normal(scale=noise_sd, size=n_points)\n",
|
||||||
|
" \n",
|
||||||
|
" # Add the noise to the sine wave\n",
|
||||||
|
" sine_wave_noise = sine_wave + noise\n",
|
||||||
|
" \n",
|
||||||
|
" # Stack the sine wave and the noisy sine wave into a 2D array\n",
|
||||||
|
" output = np.column_stack((sine_wave, sine_wave_noise))\n",
|
||||||
|
" \n",
|
||||||
|
" return output\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
"class SineDataset(Dataset):\n",
|
||||||
|
" def __init__(self, n_samples, n_points, frequency_range, phase_range, amplitude_range, noise_sd_range):\n",
|
||||||
|
" self.n_samples = n_samples\n",
|
||||||
|
" self.n_points = n_points\n",
|
||||||
|
" self.frequency_range = frequency_range\n",
|
||||||
|
" self.phase_range = phase_range\n",
|
||||||
|
" self.amplitude_range = amplitude_range\n",
|
||||||
|
" self.noise_sd_range = noise_sd_range\n",
|
||||||
|
"\n",
|
||||||
|
" def __len__(self):\n",
|
||||||
|
" return self.n_samples\n",
|
||||||
|
"\n",
|
||||||
|
" def __getitem__(self, idx):\n",
|
||||||
|
" # Generate random attributes\n",
|
||||||
|
" frequency = np.random.uniform(*self.frequency_range)\n",
|
||||||
|
" phase = np.random.uniform(*self.phase_range)\n",
|
||||||
|
" amplitude = np.random.uniform(*self.amplitude_range)\n",
|
||||||
|
" noise_sd = np.random.uniform(*self.noise_sd_range)\n",
|
||||||
|
"\n",
|
||||||
|
" # Generate sine wave with the random attributes\n",
|
||||||
|
" sine_wave = generate_sine_with_noise(self.n_points, frequency, phase, amplitude, noise_sd)\n",
|
||||||
|
"\n",
|
||||||
|
" # Return the sine wave and the parameters\n",
|
||||||
|
" return torch.Tensor(sine_wave[:-500, 1, None]), torch.Tensor(sine_wave[-500:, 0]), torch.Tensor([frequency, phase, amplitude, noise_sd])\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"\n",
|
||||||
|
"# Usage:\n",
|
||||||
|
"dataset = SineDataset(640, 1025, (1, 3), (0, 2*np.pi), (0.5, 1.5), (0.05, 1))\n",
|
||||||
|
"\n",
|
||||||
|
"def train(model, device, train_loader, optimizer, epoch):\n",
|
||||||
|
" model.train()\n",
|
||||||
|
" for batch_idx, (data, target, params) in enumerate(train_loader):\n",
|
||||||
|
" #data = data[...,None]\n",
|
||||||
|
" data, target = data.to(device), target.to(device)\n",
|
||||||
|
" optimizer.zero_grad()\n",
|
||||||
|
" output,last_state = model(data)\n",
|
||||||
|
" #import pdb;pdb.set_trace()\n",
|
||||||
|
"\n",
|
||||||
|
" loss = F.mse_loss(output, target)\n",
|
||||||
|
" loss.backward()\n",
|
||||||
|
" optimizer.step()\n",
|
||||||
|
" if batch_idx % 10 == 0:\n",
|
||||||
|
" print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
|
||||||
|
" epoch, batch_idx * len(data), len(train_loader.dataset),\n",
|
||||||
|
" 100. * batch_idx / len(train_loader), loss.item()))\n",
|
||||||
|
"\n",
|
||||||
|
"if __name__ == \"__main__\":\n",
|
||||||
|
" device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
|
||||||
|
"\n",
|
||||||
|
" model = HyenaOperatorAutoregressive1D(\n",
|
||||||
|
" d_model=128, \n",
|
||||||
|
" l_max=1024, \n",
|
||||||
|
" order=2, \n",
|
||||||
|
" filter_order=64\n",
|
||||||
|
" ).to(device)\n",
|
||||||
|
"\n",
|
||||||
|
" optimizer = optim.Adam(model.parameters())\n",
|
||||||
|
"\n",
|
||||||
|
" # Assume 10000 samples in the dataset\n",
|
||||||
|
" #dataset = SineDataset(10000, 1025, 2, 0, 1, 0.1)\n",
|
||||||
|
" train_loader = DataLoader(dataset, batch_size=64, shuffle=True)\n",
|
||||||
|
"\n",
|
||||||
|
" for epoch in range(1, 101): # Train for 10 epochs\n",
|
||||||
|
" train(model, device, train_loader, optimizer, epoch)\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "cc9f9031-5ee1-49f8-a70f-ad85ca015596",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 682,
|
||||||
|
"id": "90330622-8b44-4b45-8158-6840538f768c",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"from sklearn.linear_model import LinearRegression\n",
|
||||||
|
"from sklearn.metrics import r2_score\n",
|
||||||
|
"\n",
|
||||||
|
"def fit_and_evaluate_linear_regression(outputs_and_params):\n",
|
||||||
|
"\n",
|
||||||
|
" # Split the data into inputs (last_states) and targets (params)\n",
|
||||||
|
"\n",
|
||||||
|
" inputs = np.concatenate([x[0] for x in outputs_and_params])\n",
|
||||||
|
"\n",
|
||||||
|
" targets = np.concatenate([x[1] for x in outputs_and_params])\n",
|
||||||
|
"\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
" r2_scores = []\n",
|
||||||
|
"\n",
|
||||||
|
" param_names = [\"frequency\", \"sin_phase\", \"amplitude\", \"noise_sd\"]\n",
|
||||||
|
"\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
" # Fit the linear regression model for each parameter and calculate the R^2 score\n",
|
||||||
|
"\n",
|
||||||
|
" for i in range(targets.shape[1]):\n",
|
||||||
|
"\n",
|
||||||
|
" if param_names[i] == 'sin_phase':\n",
|
||||||
|
" # take the sine of the phase values\n",
|
||||||
|
" target_values = np.sin(targets[:, i])\n",
|
||||||
|
" \n",
|
||||||
|
" \n",
|
||||||
|
" #Not used at the moment\n",
|
||||||
|
" elif param_names[i] == 'cos_phase':\n",
|
||||||
|
" # take the cosine of the phase values\n",
|
||||||
|
" target_values = np.cos(targets[:, i])\n",
|
||||||
|
" else:\n",
|
||||||
|
" target_values = targets[:, i]\n",
|
||||||
|
"\n",
|
||||||
|
" model = LinearRegression().fit(inputs, target_values)\n",
|
||||||
|
"\n",
|
||||||
|
" pred = model.predict(inputs)\n",
|
||||||
|
"\n",
|
||||||
|
" score = r2_score(target_values, pred)\n",
|
||||||
|
"\n",
|
||||||
|
" r2_scores.append(score)\n",
|
||||||
|
"\n",
|
||||||
|
" print(f\"R^2 score for {param_names[i]}: {score:.2f}\")\n",
|
||||||
|
"\n",
|
||||||
|
" \n",
|
||||||
|
"\n",
|
||||||
|
" return r2_scores"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 683,
|
||||||
|
"id": "5eb62a22-cad8-43c4-b757-f36b6a01e9be",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def generate_outputs(model, device, data_loader):\n",
|
||||||
|
" model.eval()\n",
|
||||||
|
" outputs_and_params = []\n",
|
||||||
|
" with torch.no_grad():\n",
|
||||||
|
" for data, target, params in data_loader:\n",
|
||||||
|
" data, target = data.to(device), target.to(device)\n",
|
||||||
|
" output, last_state = model(data)\n",
|
||||||
|
" outputs_and_params.append((last_state.cpu().numpy(), params.cpu().numpy()))\n",
|
||||||
|
" return outputs_and_params\n"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 684,
|
||||||
|
"id": "a95ee542-1c39-4f04-9184-e26c6983a018",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"name": "stdout",
|
||||||
|
"output_type": "stream",
|
||||||
|
"text": [
|
||||||
|
"R^2 score for frequency: -1.09\n",
|
||||||
|
"R^2 score for sin_phase: 0.84\n",
|
||||||
|
"R^2 score for amplitude: 0.96\n",
|
||||||
|
"R^2 score for noise_sd: 0.91\n"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"outputs_and_params = generate_outputs(model, device, train_loader)\n",
|
||||||
|
"\n",
|
||||||
|
"# Fit the linear regression model and print the R^2 score for each parameter\n",
|
||||||
|
"r2_scores = fit_and_evaluate_linear_regression(outputs_and_params)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "markdown",
|
||||||
|
"id": "de65d0d2-b0c6-4ac5-a87f-70fa7f90480b",
|
||||||
|
"metadata": {},
|
||||||
|
"source": [
|
||||||
|
"# Autoregressive"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 519,
|
||||||
|
"id": "8ee139a6-aee5-4309-8685-bf0c28893279",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"def predict_autoregressive(model, initial_data, n_steps):\n",
|
||||||
|
" model.eval()\n",
|
||||||
|
" predictions = []\n",
|
||||||
|
" current_input = initial_data\n",
|
||||||
|
" with torch.no_grad():\n",
|
||||||
|
" for _ in range(n_steps):\n",
|
||||||
|
" # Get the prediction for the next step and save it\n",
|
||||||
|
" next_output, last_state = model(current_input)\n",
|
||||||
|
" predictions.append(next_output)\n",
|
||||||
|
" #import pdb;pdb.set_trace()\n",
|
||||||
|
"\n",
|
||||||
|
" # Prepare the input for the next step\n",
|
||||||
|
" next_input = torch.cat((current_input[:, 500:, :], next_output[:,:,None]), dim=1)\n",
|
||||||
|
"\n",
|
||||||
|
" current_input = next_input\n",
|
||||||
|
"\n",
|
||||||
|
" return torch.cat(predictions, dim=1)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "272040d6-4dfb-438b-a432-744e950effd1",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 594,
|
||||||
|
"id": "9f49cbf8-08e8-428a-af80-bcb2515a6327",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"torch.Size([1, 525, 1])"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 594,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"initial_data = dataset[0]\n",
|
||||||
|
"initial_data[0][None,...].shape"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 595,
|
||||||
|
"id": "c75464e3-e44d-4325-8804-29d8081e3a45",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"autoregressive_out = predict_autoregressive(model,initial_data[0][None,...],3)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 596,
|
||||||
|
"id": "90a37c56-59f3-49bc-bfbe-d9e126c42ed1",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"total = np.concatenate([initial_data[0].squeeze().numpy() ,autoregressive_out.squeeze().numpy()])"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": 597,
|
||||||
|
"id": "3d2ad0ee-7a0e-4b6f-8b40-dd565c8af84d",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"text/plain": [
|
||||||
|
"[<matplotlib.lines.Line2D at 0x7f8ec058feb0>]"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"execution_count": 597,
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "execute_result"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"data": {
|
||||||
|
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAiIAAAGhCAYAAABceN/BAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAABehElEQVR4nO3dd3xUVfo/8M+kTXoCabRA6EgPCCGoSFNA1l5YRAWWxbKwFpQV3FXssKtf3dVVZP0JuCp2hVUUpCPSS+idhARIAgTSCKlzf39M5uZOMjW5debzfr2id2Zu7j0zE+Y+c85znmMSBEEAERERkQYCtG4AERER+S8GIkRERKQZBiJERESkGQYiREREpBkGIkRERKQZBiJERESkGQYiREREpBkGIkRERKQZBiJERESkGQYiREREpBnVApF58+bBZDLhySefVOuUREREpHOqBCI7duzAggUL0Lt3bzVOR0RERAYRpPQJSktLMWHCBHz44Yd49dVXvfpdi8WCc+fOISoqCiaTSaEWEhERkZwEQUBJSQlatWqFgADXfR6KByLTpk3D2LFjMXLkSLeBSEVFBSoqKsTbZ8+eRffu3ZVuIhERESkgJycHbdq0cbmPooHIF198gd27d2PHjh0e7T937ly89NJLDe7PyclBdHS03M0jIiIiBRQXFyM5ORlRUVFu91UsEMnJycETTzyBVatWITQ01KPfmT17NmbMmCHetj2R6OhoBiJEREQG40lahUkQBEGJky9duhR33nknAgMDxftqampgMpkQEBCAiooKu8ccKS4uRkxMDIqKihiIEBERGYQ312/FekRGjBiB/fv32903efJkdOvWDc8++6zbIISIiIh8n2KBSFRUFHr27Gl3X0REBOLi4hrcT0RERP6JlVWJiIhIM4pP35Vav369mqcjIiIinWOPCBEREWmGgQgRERFphoEIERERaYaBCBEREWmGgQgRERFphoEIERERaYaBCBEREWmGgYif2XaqAJ9vz9a6GURERABULmhG2hv3n60AgI4JkRjYvrnGrSEiIn/HHhE/lX2pTOsmEBERMRAhIiIi7TAQISIiIs0wEPFTJq0bQEREBAYiREREpCEGIkRERKQZTt/1EzmXyhAeEqh1M4iIiOwwEPEDBaUVuOEf6+zuMzFJhIiIdIBDM37gaH6J1k0gIiJyiIEIERERaYaBCBEREWmGgQgRERFphoEIERERaYaBiJ/irBkiItIDBiJERESkGQYifsDElWWIiEinGIj4KQYnRESkBwxEiIiISDMMRIiIiEgzDESIiIhIMwxE/ICjqbqcvktERHrAQISIiIg0w0CEiIiINMNAhIiIiDTDQISIiIg0w0CEiIiINMNAhIiIiDSjaCAyf/589O7dG9HR0YiOjkZ6ejp+/vlnJU9JDuzMuqR1E4iIiBxSNBBp06YN5s2bh127dmHnzp0YPnw4br/9dhw8eFDJ01I9b/5yTOsmEBERORSk5MFvvfVWu9uvvfYa5s+fj61bt6JHjx5KnpqIiIgMQNFARKqmpgZff/01rly5gvT0dLVOS06YWFqViIh0QPFAZP/+/UhPT0d5eTkiIyPx/fffo3v37g73raioQEVFhXi7uLhY6eYRERGRhhSfNdO1a1dkZGRg27ZteOyxxzBx4kQcOnTI4b5z585FTEyM+JOcnKx084iIiEhDJkEQBDVPOHLkSHTs2BELFixo8JijHpHk5GQUFRUhOjpazWb6lJRZyxvc9874VNzWp5UGrSEiIl9XXFyMmJgYj67fquWI2FgsFrtgQ8psNsNsNqvcIv/EDBEiItIDRQOR2bNnY8yYMWjbti1KSkqwZMkSrF+/HitXrlTytCSzGouAi6UVSIoO1bopRETkYxQNRM6fP4+HHnoIubm5iImJQe/evbFy5UrcdNNNSp6WZDZx4XZsOnERn08dhPSOcVo3h4iIfIiigchHH32k5OFJJZtOXAQAfLrtNAMRIiKSFdea8WNf7czBwk2ZWjeDiIj8mOrJqqQff/lmHwBgVM8WaB0bpnFriIjIH7FHxE9J52xfqajWrB1EROTfGIgQERGRZhiI+CmV69gRERE5xECEWNyMiIg0w0DETz2/9IDXv8OAhYiI5MZAxE8VlzNBlYiItMdAhGBiVwcREWmEgQgRERFphoEIERERaYaBiI+yWATkFZV7uDfHZoiISBsMRHzU41/swaC5a7DyYJ7WTSEiInKKgYiP+nFfLgDgkU92adwSIiIi5xiIkMezZkycXkNERDJjIEJERESaYSDiQ9YdOY+cS2VaN4OIiMhjQVo3gOSx4dgFTF68AwCQNW+sxq0hIiLyDHtEfMTOrEuN/l1mfhARkVYYiBAREZFmGIiQx9hzQkREcmMgQkRERJphIEJERESaYSBCLFRGRESaYSDiIxhKEBGRETEQISIiIs0wECGHvSl7cwox8+u9OF9Srnp7iIjIf7CyKjlc9O72934DAFworVC5NURE5E/YI0IunThfKm4zp5WIiOTGQIRcEgStW0BERL6MgYgPen/9Ca/2twjA0bwSWCyMOoiISF0MRHzQP1Yc9Wr/uT8dxqh/bsTfVx5p8JjALhEiIlIQAxHCL4fyAQALNpzSuCVERORvGIgQERGRZhiIkEscmCEiIiUxEPEVnFtLREQGpGggMnfuXAwYMABRUVFITEzEHXfcgaNHvUukJG1Jc1UZ6hARkdwUDUQ2bNiAadOmYevWrVi1ahWqqqpw880348qVK0qelmQkcHCGiIgUpGiJ9xUrVtjdXrx4MRITE7Fr1y4MGTJEyVMTERGRAaiaI1JUVAQAaN68uZqn9QtyDpuwdojn8orK8c2uM6istmjdFCIiQ1Jt0TuLxYInn3wS1113HXr27Olwn4qKClRU1C2yVlxcrFbzqNYLyw7g1+MXxduMSRwTBAGPfroLKw9aa7CcvXwVT4zsrHGriIiMR7UekWnTpuHAgQP44osvnO4zd+5cxMTEiD/JyclqNc/wisurZDnOf7ecRuZF5vC4Ul1jQebFK2IQAgBvrz6GracKcCSvGFcrazRsHRGRsZgEFfrhp0+fjmXLlmHjxo1o37690/0c9YgkJyejqKgI0dHRSjfTUHKLruLfa09g0uAUdE6KQsqs5YqcJyHKjAsl1vfkjr6t8M/fpypyHqM4nl+Cm97e6HKfgSnN8dWj6Sq1iIhIf4qLixETE+PR9VvRoRlBEPDnP/8Z33//PdavX+8yCAEAs9kMs9msZJN8xrTPdmN3diG+230Wh18ZrXVz/Ia7IAQAtmddUqElRES+QdFAZNq0aViyZAmWLVuGqKgo5OXlAQBiYmIQFham5Kl93oGz1vyZq1XKDgPY1RHx86JpTOIlIpKfojki8+fPR1FREYYOHYqWLVuKP19++aWSp/UL6tX34MXXpsLJzJhTr9/S4D7OoiEi8oyigYggCA5/Jk2apORpiWRXYxGwdM/ZBvePuzYZAQEmXNuumd39T3+9FxXVTFolInKHa80YlElSOeSrHTkatsQ/LNl2GrO+2y/e7tMmBoM7xmHm6K4AgLfu62u3/w97z+G73Q0DFyIisqdaHRFSzl++3afYsZkWAVgsAp5fdlC8fUPneHwyJc1un7Zx4Xh7XB889eVe8b6C0goQEZFr7BExKK4Bo57l+3PF7f7tmmH+A/0d7ndnahuMvCZJvJ1XXK5424iIjI6BiEFV1TAQUcvyfXWByFePpCPS7Lwjce5dvdCmmXVG2I/7cjnThojIDQYi5BIvo0Dr2sAiyhyEwADXU5gToszisE1hWRXOXL6qePuIiIzMLwOR8yXl+GHvOWw4dkHrpuie9Bu9P1YRuXSlEh9tygQA3Ng1waPfaR8fgbDgQADADf9Yh3KFa70QERmZXwYih3NL8OfP9+DvPx/Ruim6d7lMnjVsjKi6xoJ+r6wSb/evN0XXlWtaRonbthL5RETUkF8GIrZv9hx2IFc2nyywuz35OtdLFEh9/IeB4vZn27JlaxMRka/xz0CkNhJhIiG5Il3R+MVbu3v1u1GhweL2BxtOytYmIiJf45eBSICfr5lCnpEmmg7qGNekY526UNrU5hAR+SS/DERsYYiFPSLkwsebswAAA1Oao1sL18tYu/OXb5QrOkdEZGR+GYhAHJrRthmG40cdSWWV1cgtshYku7lHkpu9HRvbu6W4XVpRLUu7iIh8jV8GIrZ1WhiHkDPnCuuGZWLCgl3s6dwdfVuL271axzS5TUREvsg/AxEmq5IbZwvryrPf1rdVo44x8ppEPDCoLQDg611ncLWS9USIiOrzz0Ck9v8MQ8iZhbVFzIZ1TYA5KLBRxzCZTHhyZBfx9je7z8jSNiIiX+KfgYjYJaJtO0iffjtxUay6mxQd2qRjxUeaxe2rlcwTISKqzy8DkYB6ccg/VhzBHxbvQI2FkQkBT36ZIW4/lJ7S5OM9NrQjAGBPdmGTj0VE5Gv8MhCxdYhkXryCXacv4f31J7H2yHn8duKitg1z4mIpS4SrSVqSvXNSZJOP1zc5FgDw84E8FPlxyXwiIkf8MhCRzkOdvGiHuF1ZbdGiMS59svU0rn11Nd5adUzrpvil4MCm/xOJlcy6Wb4/t8nHIyLyJX4ZiEgLq1bV6Hs45vmlBwAA76w5rnFLgBqLgK2nClBR7buzP86X1M2WGT+wrSzHDA8JEre/ZcIqEZEd/wxEJNsCM1Y9tizjHH7/n62Y/d1+rZuimDdWHBW3Hx/RSZZjSiv47jp9WZZjEhH5Cv8MRCRdInovJaLHZXG+231W6yYopvBqXQ5HpDnIxZ6eu6Zl08rDExH5Mr8MRAIkF3edxyGkspDanJCgAJPdCrpNOmZQAL59bLB4W5oMS0Tk7/wyEDHBcSSix94HHTZJl574Yg/uW7AFliZOwbbNUHprXF8ZWlVHOvvmvXUnZD22L8m6eAV7sjl8ReRP5Ol7NhiTXRyi7z4Rk8mk//EjHViWcQ4Line truncated
|
||||||
|
"text/plain": [
|
||||||
|
"<Figure size 640x480 with 1 Axes>"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"metadata": {},
|
||||||
|
"output_type": "display_data"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"source": [
|
||||||
|
"import matplotlib.pyplot as plt\n",
|
||||||
|
"\n",
|
||||||
|
"plt.plot(total)"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "eb368900-34b3-4c36-a916-62ed3a97173b",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "50253cc9-254c-4933-94bd-36c28d0f97b0",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"id": "895a5767-2865-477f-aa21-094a711bf033",
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": []
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": "Python 3 (ipykernel)",
|
||||||
|
"language": "python",
|
||||||
|
"name": "python3"
|
||||||
|
},
|
||||||
|
"language_info": {
|
||||||
|
"codemirror_mode": {
|
||||||
|
"name": "ipython",
|
||||||
|
"version": 3
|
||||||
|
},
|
||||||
|
"file_extension": ".py",
|
||||||
|
"mimetype": "text/x-python",
|
||||||
|
"name": "python",
|
||||||
|
"nbconvert_exporter": "python",
|
||||||
|
"pygments_lexer": "ipython3",
|
||||||
|
"version": "3.10.10"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 5
|
||||||
|
}
|
||||||
Reference in new issue
Block a user