fix: fixed model shapes + redit training loop
This commit is contained in:
parent
ed44d5b283
commit
eb4a014aa1
3 changed files with 68 additions and 56 deletions
|
|
@ -12,19 +12,13 @@ from train import train
|
|||
|
||||
|
||||
def create_model(trial: tr.Trial, vocab_size: int = 256):
|
||||
num_layers = trial.suggest_int("num_layers", 1, 6)
|
||||
hidden_dim = trial.suggest_int("hidden_dim", 64, 512, log=True)
|
||||
kernel_size = trial.suggest_int("kernel_size", 2, 7)
|
||||
dropout_prob = trial.suggest_float("dropout_prob", 0.1, 0.5)
|
||||
use_batchnorm = trial.suggest_categorical("use_batchnorm", [True, False])
|
||||
embedding_dim = trial.suggest_int("embedding_dim", 64, 512, log=True)
|
||||
|
||||
return CNNPredictor(
|
||||
vocab_size=vocab_size,
|
||||
num_layers=num_layers,
|
||||
hidden_dim=hidden_dim,
|
||||
kernel_size=kernel_size,
|
||||
dropout_prob=dropout_prob,
|
||||
use_batchnorm=use_batchnorm
|
||||
embed_dim=embedding_dim,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
Reference in a new issue