File size: 5,819 Bytes
22a6d48 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"extend_to = None #<= CHANGE THIS. The original is 178 symbols\n",
"\n",
"save_path = \"./Extend/New_Weights\"\n",
"config_path = \"./Configs/config.yaml\"\n",
"model_path = \"./Models/Finetune/base_model.pth\"\n",
"\n",
"#⚠️ Must run this notebook first before adding any symbol to the config file"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# load packages\n",
"%cd ..\n",
"import yaml\n",
"import torch\n",
"from torch import nn\n",
"import os\n",
"from models import *\n",
"from utils import *\n",
"import warnings\n",
"warnings.filterwarnings(\"ignore\")\n",
"\n",
"device = 'cpu'\n",
"\n",
"config = yaml.safe_load(open(config_path, \"r\", encoding=\"utf-8\"))\n",
"try:\n",
" symbols = (\n",
" list(config['symbol']['pad']) +\n",
" list(config['symbol']['punctuation']) +\n",
" list(config['symbol']['letters']) +\n",
" list(config['symbol']['letters_ipa']) +\n",
" list(config['symbol']['extend'])\n",
" )\n",
" symbol_dict = {}\n",
" for i in range(len((symbols))):\n",
" symbol_dict[symbols[i]] = i\n",
"\n",
" n_token = len(symbol_dict) + 1\n",
" print(\"\\nFound\", n_token, \"symbols in the original config file\")\n",
"except Exception as e:\n",
" print(f\"\\nERROR: Cannot find {e} in config file!\\nYour config file is likely outdated, please download updated version from the repository.\")\n",
" raise SystemExit(1)\n",
"\n",
"if (extend_to-n_token) <= 0:\n",
" print(f\"\\nERROR: Cannot extend from {n_token} to {extend_to}.\")\n",
" raise SystemExit(1)\n",
"\n",
"model_params = recursive_munch(config['model_params'])\n",
"model_params['n_token'] = n_token\n",
"model = build_model(model_params)\n",
"\n",
"keys_to_keep = {'predictor', 'decoder', 'text_encoder', 'style_encoder', 'text_aligner', 'pitch_extractor', 'mpd', 'msd'}\n",
"params_whole = torch.load(model_path, map_location='cpu')\n",
"params = params_whole['net']\n",
"params = {key: value for key, value in params.items() if key in keys_to_keep}\n",
"\n",
"for key in list(model.keys()):\n",
" if key not in keys_to_keep:\n",
" del model[key]\n",
"\n",
"for key in model:\n",
" if key in params:\n",
" print('%s loaded' % key)\n",
" try:\n",
" model[key].load_state_dict(params[key])\n",
" except:\n",
" from collections import OrderedDict\n",
" state_dict = params[key]\n",
" new_state_dict = OrderedDict()\n",
" for k, v in state_dict.items():\n",
" name = k[7:] # remove `module.`\n",
" new_state_dict[name] = v\n",
" # load params\n",
" model[key].load_state_dict(new_state_dict, strict=False)\n",
"\n",
"old_weight = [\n",
" model['text_encoder'].embedding,\n",
" model['text_aligner'].ctc_linear[2].linear_layer,\n",
" model['text_aligner'].asr_s2s.embedding,\n",
" model['text_aligner'].asr_s2s.project_to_n_symbols\n",
"]\n",
"print(\"\\nOld shape:\") \n",
"for module in old_weight:\n",
" print(module, module.weight.shape)\n",
"\n",
"for i in range(len(old_weight)):\n",
" new_shape = (extend_to, old_weight[i].weight.shape[1])\n",
" new_weight = torch.randn(new_shape) * 0.01 #init mean=0, std=0.01\n",
" with torch.no_grad():\n",
" new_weight[:old_weight[i].weight.size(0), :] = old_weight[i].weight.detach().clone()\n",
" new_param = nn.Parameter(new_weight, requires_grad=True)\n",
"\n",
" if isinstance(old_weight[i], nn.Embedding):\n",
" old_weight[i].num_embeddings = extend_to\n",
" \n",
" if isinstance(old_weight[i], nn.Linear):\n",
" old_weight[i].out_features = extend_to\n",
" #update bias\n",
" old_bias = old_weight[i].bias.detach()\n",
" old_dim = old_bias.shape[0]\n",
" new_bias = torch.zeros(extend_to)\n",
" new_bias[:old_dim] = old_bias.clone()\n",
" old_weight[i].bias.data = new_bias\n",
"\n",
" old_weight[i].weight = new_param\n",
"\n",
"print(\"\\nNew shape:\")\n",
"for module in old_weight:\n",
" print(module, module.weight.shape)\n",
"\n",
"if not os.path.exists(save_path):\n",
" os.mkdir(save_path)\n",
"\n",
"print(f\"\\n\\n✅ Successfully extended the token set to a maximum of {extend_to} symbols.\")\n",
"print(f\"You can now add {extend_to - n_token} additional symbols in the config file.\")\n",
"\n",
"#save new weights\n",
"state = {\n",
" 'net': {key: model[key].state_dict() for key in model}, \n",
" 'optimizer': None,\n",
" 'iters': 0,\n",
" 'val_loss': 0,\n",
" 'epoch': 0,\n",
"}\n",
"torch.save(state, os.path.join(save_path, 'extended.pth'))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "base",
"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.11.7"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
|