How to specify training arguments for huggingface transformer using skorch

username: ichrist97
cooked: <p>I am trying to fine tune a huggingface transformer using <code>skorch</code>. I followed the example notebook from skorch for the implementation (<a href="https://nbviewer.org/github/skorch-dev/skorch/blob/master/notebooks/Hugging_Face_Finetuning.ipynb" class="inline-onebox" rel="noopener nofollow ugc">Jupyter Notebook Viewer</a>)</p>
<p>The fine tuning works like in the example notebook, but now I want to apply <code>RandomizedSearchCV</code> from sklearn to tune the hyperparameters of the transformer model.</p>
<p>I am able to pass the parameter ranges to the pytorch module in my pipeline (<code>BertModule</code>), but neither the skorch documentation nor the huggingface documentation specifies how the hyperparameters can be used for training in the pytorch module.</p>
<p>Normally, with the Huggingface Trainer API some <code>TrainingArguments</code>could be defined, but I don’t see how they can be used with skorch.</p>
<p>In this case, I want to tune the hyper parameters <code>weight_decay</code> and <code>warmup_steps</code> for the transformer model.</p>
<p>My code looks like the following:</p>
<pre><code class="lang-python">class BertModule(nn.Module):
    def __init__(
        self,
        name: str,
        num_labels: int,
        # where to use them though?
        weight_decay: float = 0.0,
        warmup_steps: int = 0,
    ):
        super().__init__()
        self.name = name
        self.num_labels = num_labels

        self.reset_weights()

    def reset_weights(self):
        self.bert = AutoModelForSequenceClassification.from_pretrained(
            self.name, num_labels=self.num_labels
        )

    def forward(self, **kwargs):
        pred = self.bert(**kwargs)
        return pred.logits

pipe = Pipeline(
    [
        ("tokenizer", HuggingfacePretrainedTokenizer(TOKENIZER)),
        (
            "net",
            NeuralNetClassifier(
                BertModule,
                module__name=PRETRAINED_MODEL,
                module__num_labels=len(set(y_train)),
                module__weight_decay=0.0,
                module__warmup_steps=0,
                optimizer=OPTIMIZER,
                lr=LR,
                max_epochs=MAX_EPOCHS,
                criterion=CRITERION,
                batch_size=BATCH_SIZE,
                iterator_train__shuffle=True,
                device=DEVICE,
                callbacks=[
                    LRScheduler(LambdaLR, lr_lambda=lr_schedule, step_every="batch"),
                    ProgressBar(),
                ],
            ),
        ),
    ]
)

params = {
    "net__lr": loguniform(1e-5, 1e-3),
    "net__max_epochs": randint(2, 5),
    "net__batch_size": [8, 16, 32, 64, 128],
    "net__module__weight_decay": uniform(0, 0.3),
    "net__module__warmup_steps": randint(0, 1000),
}

rs = RandomizedSearchCV(
    pipe, params, n_iter=10, verbose=2, refit=True, scoring="accuracy", cv=5
)

rs.fit(X_train, y_train)

</code></pre>
<p>Thanks for the help! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>