Tutorial_Training_for_BO.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f32e2ce2",
   "metadata": {},
   "source": [
    "## Train your own models\n",
    "To train you simply need to call `train.train`.\n",
    "We give all necessary code. The most important bits are in the `priors` dir, e.g. `hebo_prior`, it stores the priors\n",
    "with which we train our models.\n",
    "\n",
    "### Training the HEBO+ model, `model_hebo_morebudget_9_unused_features_3.pt`\n",
    "You can train this model on 8 GPUs using `torchrun` or `submitit`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4949ecba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from pfns import priors, encoders, utils, bar_distribution, train\n",
    "from ConfigSpace import hyperparameters as CSH\n",
    "\n",
    "config_heboplus = {\n",
    "     'priordataloader_class_or_get_batch': priors.get_batch_to_dataloader(\n",
    "         priors.get_batch_sequence(\n",
    "             priors.hebo_prior.get_batch,\n",
    "             priors.utils.sample_num_feaetures_get_batch,\n",
    "         )\n",
    "     ),\n",
    "     'encoder_generator': encoders.get_normalized_uniform_encoder(encoders.get_variable_num_features_encoder(encoders.Linear)),\n",
    "     'emsize': 512,\n",
    "     'nhead': 4,\n",
    "     'warmup_epochs': 5,\n",
    "     'y_encoder_generator': encoders.Linear,\n",
    "     'batch_size': 128,\n",
    "     'scheduler': utils.get_cosine_schedule_with_warmup,\n",
    "     'extra_prior_kwargs_dict': {'num_features': 18,\n",
    "      'hyperparameters': {\n",
    "       'lengthscale_concentration': 1.2106559584074301,\n",
    "       'lengthscale_rate': 1.5212245992840594,\n",
    "       'outputscale_concentration': 0.8452312502679863,\n",
    "       'outputscale_rate': 0.3993553245745406,\n",
    "       'add_linear_kernel': False,\n",
    "       'power_normalization': False,\n",
    "       'hebo_warping': False,\n",
    "       'unused_feature_likelihood': 0.3,\n",
    "       'observation_noise': True}},\n",
    "     'epochs': 50,\n",
    "     'lr': 0.0001,\n",
    "     'seq_len': 60,\n",
    "     'single_eval_pos_gen': utils.get_uniform_single_eval_pos_sampler(50, min_len=1), #<function utils.get_uniform_single_eval_pos_sampler.<locals>.<lambda>()>,\n",
    "     'aggregate_k_gradients': 2,\n",
    "     'nhid': 1024,\n",
    "     'steps_per_epoch': 1024,\n",
    "     'weight_decay': 0.0,\n",
    "     'train_mixed_precision': False,\n",
    "     'efficient_eval_masking': True,\n",
    "     'nlayers': 12}\n",
    "\n",
    "\n",
    "config_heboplus_userpriors = {**config_heboplus,\n",
    "    'priordataloader_class_or_get_batch': priors.get_batch_to_dataloader(\n",
    "                              priors.get_batch_sequence(\n",
    "                                  priors.hebo_prior.get_batch,\n",
    "                                  priors.condition_on_area_of_opt.get_batch,\n",
    "                                  priors.utils.sample_num_feaetures_get_batch\n",
    "                              )),\n",
    "    'style_encoder_generator': encoders.get_normalized_uniform_encoder(encoders.get_variable_num_features_encoder(encoders.Linear))\n",
    "}\n",
    "\n",
    "config_bnn = {'priordataloader_class_or_get_batch': priors.get_batch_to_dataloader(\n",
    "         priors.get_batch_sequence(\n",
    "             priors.simple_mlp.get_batch,\n",
    "             priors.input_warping.get_batch,\n",
    "             priors.utils.sample_num_feaetures_get_batch,\n",
    "         )\n",
    "     ),\n",
    "     'encoder_generator': encoders.get_normalized_uniform_encoder(encoders.get_variable_num_features_encoder(encoders.Linear)),\n",
    "     'emsize': 512,\n",
    "     'nhead': 4,\n",
    "     'warmup_epochs': 5,\n",
    "     'y_encoder_generator': encoders.Linear,\n",
    "     'batch_size': 128,\n",
    "     'scheduler': utils.get_cosine_schedule_with_warmup,\n",
    "     'extra_prior_kwargs_dict': {'num_features': 18,\n",
    "      'hyperparameters': {'mlp_num_layers': CSH.UniformIntegerHyperparameter('mlp_num_layers', 8, 15),\n",
    "       'mlp_num_hidden': CSH.UniformIntegerHyperparameter('mlp_num_hidden', 36, 150),\n",
    "       'mlp_init_std': CSH.UniformFloatHyperparameter('mlp_init_std',0.08896049884896237, 0.1928554813280186),\n",
    "       'mlp_sparseness': 0.1449806273312999,\n",
    "       'mlp_input_sampling': 'uniform',\n",
    "       'mlp_output_noise': CSH.UniformFloatHyperparameter('mlp_output_noise', 0.00035983014290491186, 0.0013416342770574585),\n",
    "       'mlp_noisy_targets': True,\n",
    "       'mlp_preactivation_noise_std': CSH.UniformFloatHyperparameter('mlp_preactivation_noise_std',0.0003145707276259681, 0.0013753183831259406),\n",
    "       'input_warping_c1_std': 0.9759720822120248,\n",
    "       'input_warping_c0_std': 0.8002534583197192,\n",
    "       'num_hyperparameter_samples_per_batch': 16}\n",
    "                                 },\n",
    "     'epochs': 50,\n",
    "     'lr': 0.0001,\n",
    "     'seq_len': 60,\n",
    "     'single_eval_pos_gen': utils.get_uniform_single_eval_pos_sampler(50, min_len=1), \n",
    "     'aggregate_k_gradients': 1,\n",
    "     'nhid': 1024,\n",
    "     'steps_per_epoch': 1024,\n",
    "     'weight_decay': 0.0,\n",
    "     'train_mixed_precision': True,\n",
    "     'efficient_eval_masking': True,\n",
    "}\n",
    "\n",
    "\n",
    "# now let's add the criterions, where we decide the border positions based on the prior\n",
    "def get_ys(config, device='cuda:0'):\n",
    "    bs = 128\n",
    "    all_targets = []\n",
    "    for num_hps in [2,8,12]: # a few different samples in case the number of features makes a difference in y dist\n",
    "        b = config['priordataloader_class_or_get_batch'].get_batch_method(\n",
    "            bs,1000,  num_hps, epoch=0, device=device, hyperparameters={**config['extra_prior_kwargs_dict']['hyperparameters'],\n",
    "                                                                        'num_hyperparameter_samples_per_batch': -1,})\n",
    "        all_targets.append(b.target_y.flatten())\n",
    "    return torch.cat(all_targets,0)\n",
    "\n",
    "def add_criterion(config, device='cuda:0'):\n",
    "    return {**config, 'criterion': bar_distribution.FullSupportBarDistribution(\n",
    "        bar_distribution.get_bucket_limits(1000,ys=get_ys(config,device).cpu())\n",
    "    )}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c8789428",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore', category=UserWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4dec5326",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 384000 y evals to estimate 1000 buckets. Cut off the last 0 ys.\n",
      "Using cpu:0 device\n",
      "init dist\n",
      "Not using distributed\n",
      "DataLoader.__dict__ {'num_steps': 1024, 'get_batch_kwargs': {'batch_size': 128, 'eval_pos_seq_len_sampler': <function train.<locals>.eval_pos_seq_len_sampler at 0x7f817c751ee0>, 'seq_len_maximum': 60, 'device': 'cpu:0', 'num_features': 18, 'hyperparameters': {'lengthscale_concentration': 1.2106559584074301, 'lengthscale_rate': 1.5212245992840594, 'outputscale_concentration': 0.8452312502679863, 'outputscale_rate': 0.3993553245745406, 'add_linear_kernel': False, 'power_normalization': False, 'hebo_warping': False, 'unused_feature_likelihood': 0.3, 'observation_noise': True}}, 'num_features': 18, 'epoch_count': 0}\n",
      "Style definition of first 3 examples: None\n",
      "Initialized decoder for standard with (None, 1000)  and nout 1000\n",
      "Using a Transformer with 26.79 M parameters\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(inf,\n",
       " inf,\n",
       " TransformerModel(\n",
       "   (transformer_encoder): TransformerEncoderDiffInit(\n",
       "     (layers): ModuleList(\n",
       "       (0): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (1): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (2): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (3): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (4): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (5): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (6): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (7): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (8): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (9): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (10): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "       (11): TransformerEncoderLayer(\n",
       "         (self_attn): MultiheadAttention(\n",
       "           (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
       "         )\n",
       "         (linear1): Linear(in_features=512, out_features=1024, bias=True)\n",
       "         (dropout): Dropout(p=0.0, inplace=False)\n",
       "         (linear2): Linear(in_features=1024, out_features=512, bias=True)\n",
       "         (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "         (dropout1): Dropout(p=0.0, inplace=False)\n",
       "         (dropout2): Dropout(p=0.0, inplace=False)\n",
       "       )\n",
       "     )\n",
       "   )\n",
       "   (encoder): Sequential(\n",
       "     (0): Normalize()\n",
       "     (1): VariableNumFeaturesEncoder(\n",
       "       (base_encoder): Linear(in_features=18, out_features=512, bias=True)\n",
       "     )\n",
       "   )\n",
       "   (y_encoder): Linear(in_features=1, out_features=512, bias=True)\n",
       "   (pos_encoder): NoPositionalEncoding()\n",
       "   (decoder_dict): ModuleDict(\n",
       "     (standard): Sequential(\n",
       "       (0): Linear(in_features=512, out_features=1024, bias=True)\n",
       "       (1): GELU(approximate='none')\n",
       "       (2): Linear(in_features=1024, out_features=1000, bias=True)\n",
       "     )\n",
       "   )\n",
       "   (criterion): FullSupportBarDistribution()\n",
       " ),\n",
       " <pfns.priors.utils.get_batch_to_dataloader.<locals>.DL at 0x7f809d906520>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Now let's train either with\n",
    "train.train(**add_criterion(config_heboplus,device='cpu:0'))\n",
    "# or\n",
    "#train.train(**add_criterion(config_heboplus_userpriors))\n",
    "# or\n",
    "#train.train(**add_criterion(config_bnn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c427f1e1",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}