fomaml_step#
- ivy.fomaml_step(batch, inner_cost_fn, outer_cost_fn, variables, inner_grad_steps, inner_learning_rate, /, *, inner_optimization_step=<function gradient_descent_update>, inner_batch_fn=None, outer_batch_fn=None, average_across_steps=False, batched=True, inner_v=None, keep_inner_v=True, outer_v=None, keep_outer_v=True, return_inner_v=False, num_tasks=None, stop_gradients=True)[source]#
Perform step of first order MAML.
- Parameters:
batch (
Container
) – The input batchinner_cost_fn (
Callable
) – callable for the inner loop cost function, receiving sub-batch, inner vars and outer varsouter_cost_fn (
Callable
) – callable for the outer loop cost function, receiving task-specific sub-batch, inner vars and outer vars. If None, the cost from the inner loop will also be optimized in the outer loop.variables (
Container
) – Variables to be optimized during the meta stepinner_grad_steps (
int
) – Number of gradient steps to perform during the inner loop.inner_learning_rate (
float
) – The learning rate of the inner loop.inner_optimization_step (
Callable
) – The function used for the inner loop optimization. (default:<function gradient_descent_update at 0x7f9661d19310>
) Default is ivy.gradient_descent_update.inner_batch_fn (
Optional
[Callable
]) – Function to apply to the task sub-batch, before passing to the inner_cost_fn. (default:None
) Default isNone
.outer_batch_fn (
Optional
[Callable
]) – Function to apply to the task sub-batch, before passing to the outer_cost_fn. (default:None
) Default isNone
.average_across_steps (
bool
) – Whether to average the inner loop steps for the outer loop update. (default:False
) Default isFalse
.batched (
bool
) – Whether to batch along the time dimension, and run the meta steps in batch. (default:True
) Default isTrue
.inner_v (
Optional
[Container
]) – Nested variable keys to be optimized during the inner loop, with same keys and (default:None
) boolean values. (Default value = None)keep_inner_v (
bool
) – If True, the key chains in inner_v will be kept, otherwise they will be removed. (default:True
) Default isTrue
.outer_v (
Optional
[Container
]) – Nested variable keys to be optimized during the inner loop, with same keys and (default:None
) boolean values. (Default value = None)keep_outer_v (
bool
) – If True, the key chains in inner_v will be kept, otherwise they will be removed. (default:True
) Default isTrue
.return_inner_v (
Union
[str
,bool
]) – Either ‘first’, ‘all’, or False. ‘first’ means the variables for the first task (default:False
) inner loop will also be returned. variables for all tasks will be returned with ‘all’. Default isFalse
.num_tasks (
Optional
[int
]) – Number of unique tasks to inner-loop optimize for the meta step. Determined from (default:None
) batch by default.stop_gradients (
bool
) – Whether to stop the gradients of the cost. (default:True
) Default isTrue
.
- Return type:
- Returns:
ret – The cost and the gradients with respect to the outer loop variables.