164 lines
3.4 KiB
Plaintext
164 lines
3.4 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"%load_ext autoreload\n",
|
|
"%autoreload 2\n",
|
|
"\n",
|
|
"import data\n",
|
|
"import show\n",
|
|
"import model"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"x,y = data.detaset(data.load())\n",
|
|
"len(x)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "2",
|
|
"metadata": {
|
|
"editable": true,
|
|
"slideshow": {
|
|
"slide_type": ""
|
|
},
|
|
"tags": []
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"sample=x[420]\n",
|
|
"show.showSpec(sample)\n",
|
|
"show.playSpec(sample)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from tinygrad import nn\n",
|
|
"gen = model.Gen()\n",
|
|
"optimizer = nn.optim.AdamW(nn.state.get_parameters(gen), lr=1e-4)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"@TinyJit\n",
|
|
"def jit_step(X: Tensor, Y: Tensor,epoch) -> Tensor:\n",
|
|
" Tensor.training = True\n",
|
|
" optimizer.zero_grad()\n",
|
|
" _, loss = gen.__Lcall__(X,Y,epoch)\n",
|
|
" loss.backward()\n",
|
|
" optimizer.step()\n",
|
|
" return loss.realize()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "5",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import numpy as np\n",
|
|
"x_np, y_np = np.stack(x), np.stack(y)\n",
|
|
"x_np = np.expand_dims(x_np, axis=1)\n",
|
|
"y_np = np.expand_dims(y_np, axis=1)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6",
|
|
"metadata": {
|
|
"editable": true,
|
|
"slideshow": {
|
|
"slide_type": ""
|
|
},
|
|
"tags": []
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"size=len(x)\n",
|
|
"BACH_SIZE=32\n",
|
|
"eshape = (BACH_SIZE, 1, 128, 216)\n",
|
|
"\n",
|
|
"for e in range(0,12):\n",
|
|
" print(f\"\\n--- Starting Epoch {e} ---\\n\")\n",
|
|
" l=0\n",
|
|
" \n",
|
|
" permutation = np.random.permutation(size)\n",
|
|
" x_np = x_np[permutation]\n",
|
|
" y_np = y_np[permutation]\n",
|
|
" \n",
|
|
" for i in range(0,size,BACH_SIZE):\n",
|
|
" tx,ty=Tensor(x_np[i:i+BACH_SIZE]),Tensor(y_np[i:i+BACH_SIZE])\n",
|
|
" if(tx.shape != eshape or ty.shape != eshape):\n",
|
|
" continue\n",
|
|
" l+=jit_step(tx,ty,e).numpy()\n",
|
|
" \n",
|
|
" l /= (size/BACH_SIZE)\n",
|
|
" print(f\"loss of {l}\")\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "7",
|
|
"metadata": {
|
|
"editable": true,
|
|
"slideshow": {
|
|
"slide_type": ""
|
|
},
|
|
"tags": []
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"from tinygrad.nn.state import safe_save, get_state_dict\n",
|
|
"safe_save(get_state_dict(gen),\"music.safetensors\")"
|
|
]
|
|
}
|
|
],
|
|
"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.12.11"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|