catenets.models.torch.slearner module
- class SLearner(n_unit_in: int, binary_y: bool, po_estimator: Optional[Any] = None, n_layers_out: int = 2, n_units_out: int = 100, n_units_out_prop: int = 100, n_layers_out_prop: int = 2, weight_decay: float = 0.0001, lr: float = 0.0001, n_iter: int = 10000, batch_size: int = 100, val_split_prop: float = 0.3, n_iter_print: int = 50, seed: int = 42, nonlin: str = 'elu', weighting_strategy: Optional[str] = None, batch_norm: bool = True, early_stopping: bool = True, dropout: bool = False, dropout_prob: float = 0.2)
Bases:
catenets.models.torch.base.BaseCATEEstimator
S-learner for treatment effect estimation (single learner, treatment indicator just another feature).
- Parameters
n_unit_in (int) – Number of features
binary_y (bool) – Whether the outcome is binary
po_estimator (sklearn/PyTorch model, default: None) – Custom potential outcome model. If this parameter is set, the rest of the parameters are ignored.
n_layers_out (int) – Number of hypothesis layers (n_layers_out x n_units_out + 1 x Linear layer)
n_layers_out_prop (int) – Number of hypothesis layers for propensity score(n_layers_out x n_units_out + 1 x Linear layer)
n_units_out (int) – Number of hidden units in each hypothesis layer
n_units_out_prop (int) – Number of hidden units in each propensity score hypothesis layer
weight_decay (float) – l2 (ridge) penalty
lr (float) – learning rate for optimizer
n_iter (int) – Maximum number of iterations
batch_size (int) – Batch size
val_split_prop (float) – Proportion of samples used for validation split (can be 0)
n_iter_print (int) – Number of iterations after which to print updates
seed (int) – Seed used
nonlin (string, default 'elu') – Nonlinearity to use in the neural net. Can be ‘elu’, ‘relu’, ‘selu’ or ‘leaky_relu’.
weighting_strategy (optional str, None) – Whether to include propensity head and which weightening strategy to use
- _backward_hooks: Dict[int, Callable]
- _buffers: Dict[str, Optional[torch.Tensor]]
- _create_extended_matrices(X: torch.Tensor) torch.Tensor
- _forward_hooks: Dict[int, Callable]
- _forward_pre_hooks: Dict[int, Callable]
- _is_full_backward_hook: Optional[bool]
- _load_state_dict_post_hooks: Dict[int, Callable]
- _load_state_dict_pre_hooks: Dict[int, Callable]
- _modules: Dict[str, Optional[Module]]
- _non_persistent_buffers_set: Set[str]
- _parameters: Dict[str, Optional[torch.nn.parameter.Parameter]]
- _state_dict_hooks: Dict[int, Callable]
- fit(X: torch.Tensor, y: torch.Tensor, w: torch.Tensor) catenets.models.torch.slearner.SLearner
Fit treatment models.
- Parameters
X (torch.Tensor of shape (n_samples, n_features)) – The features to fit to
y (torch.Tensor of shape (n_samples,) or (n_samples, )) – The outcome variable
w (torch.Tensor of shape (n_samples,)) – The treatment indicator
- predict(X: torch.Tensor, return_po: bool = False, training: bool = False) torch.Tensor
Predict treatment effects and potential outcomes
- Parameters
X (array-like of shape (n_samples, n_features)) – Test-sample features
- Returns
y
- Return type
array-like of shape (n_samples,)
- training: bool