Training API

The narrative guides for these classes are the Training Overview and the training subpages.

RelationalExample

class relational_transformers.RelationalExample(input, label, target=None)

One labeled context used for training or evaluation.

target is required when input is a raw cell-vector array. Typed RelationalBatch inputs already carry their target mask and leave it unset.

Parameters:

  • input (Any)

  • label (Any)

  • target (int | Sequence[int] | None)

RelationalTrainingArguments

class relational_transformers.RelationalTrainingArguments(output_dir: 'str' = 'relational_model', num_train_epochs: 'int' = 1, per_device_train_batch_size: 'int' = 8, learning_rate: 'float' = 1e-05, weight_decay: 'float' = 0.01, max_grad_norm: 'float' = 1.0, gradient_accumulation_steps: 'int' = 1, seed: 'int' = 42, logging_steps: 'int' = 10, save_strategy: 'str' = 'epoch', training_backend: 'str' = 'torch')
Parameters:

  • output_dir (str)

  • num_train_epochs (int)

  • per_device_train_batch_size (int)

  • learning_rate (float)

  • weight_decay (float)

  • max_grad_norm (float)

  • gradient_accumulation_steps (int)

  • seed (int)

  • logging_steps (int)

  • save_strategy (str)

  • training_backend (str)

RelationalTrainer

class relational_transformers.RelationalTrainer(*, model, args, train_dataset, task=None, problem_type=None)

Small, dependency-free trainer for complete RT-J fine-tuning.

Parameters:

TaskHead

class relational_transformers.TaskHead(d_model, num_labels=1, problem_type='binary')

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

  • d_model (int)

  • num_labels (int)

  • problem_type (str)

forward(features)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

features (Tensor)

Return type:

Tensor

fit_head

relational_transformers.training.fit_head(transformer, examples, *, task, num_labels=1, problem_type='binary', epochs=100, learning_rate=0.001, weight_decay=0.0001)

Encode each example once and fit a lightweight task head.

Parameters:

  • examples (Sequence[RelationalExample])

  • task (str)

  • num_labels (int)

  • problem_type (str)

  • epochs (int)

  • learning_rate (float)

  • weight_decay (float)

Return type:

TaskHead