<filename>baselines/turtle/turtle_code.py
import torch
import numpy as np

from algorithm import Algorithm
from shell import MetaLearner, get_init_info, get_fast_weights,\
                           INP_TO_CODE, HIST_TO_CODE, Input, History
from utils import put_on_device, get_init_score_and_operator, accuracy, get_params




class Turtle(Algorithm):
    """sTateless neURal meTa LEarning (TURTLE)
    
    Meta-learning algorithm that combines the best of MAML and 
    LSTM meta-learners. It learns both a good common initialization 
    point for tasks, and learns how to make parameter updates in a
    using a meta-learner network.
    
    ...

    Attributes
    ----------
    baselearner_fn : constructor function
        Constructor function for the base-learner
    baselearner_args : dict
        Dictionary of keyword arguments for the base-learner
    opt_fn : constructor function
        Constructor function for the optimizer to use
    T : int
        Number of update steps to parameters per task
    train_batch_size : int
        Indicating the size of minibatches that are sampled from meta-train tasks
    test_batch_size : int
        Size of batches to sample from meta-[val/test] tasks
    lr : float
        Learning rate for the optimizer
    validation : boolean
        Whether this model should use meta-validation
    dev : str
        Device identifier
    trainable : boolean
        Whether it can train on meta-trrain tasks
    episodic : boolean
        Whether to sample tasks or mini batches for training
        
    Methods
    -------
    train(train_x, train_y, test_x, test_y)
        Perform a single training step on a given task
    
    evaluate(train_x, train_y, test_x, test_y)
        Evaluate the performance on the given task
        
    dump_state()
        Dump the meta-learner state s.t. it can be loaded again later
        
    load_state(state)
        Set meta-learner state to provided @state 
    """
    
    def __init__(self, input_type, history, beta, grad_clip, layers, act, 
                 second_order, layer_wise, param_lr, decouple, meta_batch_size=1, time_input=False, **kwargs):
        """Set all (inherited) attributes
        
        Attributes
        ----------
        input_type : str
            Specifier of the input type (raw_grads, raw_loss_grads, proc_loss_grads,
            maml). Determines the input to the meta-learner network
        layers : list
            List of integers corresponding to the number of neurons per hidden/output layer
            e.g., [5,5,1] is a neural network [input -> hidden 5 -> hidden 5 -> output 1]
        history : str
            Type of historical information to use [none, grads, updates]
        beta : float
            Float between [0,1] that determines the memory capacity for the historical
            information. Only affects results if history != none
        grad_clip : float
            Gradient clipping value to use
        act : act_fn
            Activation function to user for meta-learner
        second_order : boolean
            Whether to use second-order gradient information
        layer_wise : boolean
            Whether to use a meta-learner network for every individual base-learner layer
        meta_batch_size : int
            Number of tasks to compute outer-update
        time_input : boolean
            Whether to add a timestamp to the input (this way TURTLE can find an equivalent of a learning schedule)
        param_lr : boolean
            Whether to maintain a learning rate per parameter
        decouple : int
            Whether to decouple the base- and meta-search processes (int indicates the task # to switch modes)
        **kwargs : dict
            Ignored keyword arguments
        """
        
        super().__init__(**kwargs)
        self.input_type = input_type
        self.input_code = INP_TO_CODE[input_type]
        self.history_code = HIST_TO_CODE[history]
        self.beta = beta
        self.grad_clip = grad_clip
        self.second_order = second_order
        self.layer_wise = layer_wise
        self.meta_batch_size = meta_batch_size
        self.time_input = time_input
        self.param_lr = param_lr
        self.decouple = decouple
        self.metasearch = bool(self.decouple) # currently looking only for optimizer if self.decouple
        self.dcounter = 0
        # Increment after every train step on a single task, and update
        # init when task_counter % meta_batch_size == 0
        self.task_counter = 0 

        # Get random initialization point for baselearner
        self.baselearner = self.baselearner_fn(**self.baselearner_args).to(self.dev)
        self.initialization, self.slices = get_params(self.baselearner, self.dev)

        # If we do not use layer-wise meta-learner networks, there is only a single parameter slice
        # namely (0, num_params)
        if not self.layer_wise:
            self.slices = [ tuple([0, self.slices[-1][1]]) ]
            print(self.slices) 

        self.num_params = 0
        # Enable gradient tracking for the initialization parameters    
        for p in self.initialization:
            self.num_params += p.numel()
            p.requires_grad = True

        # We do not use historical information when T = 1
        if self.T == 1:
            ml_hist = "none" 
        else:
            ml_hist = history

        # Create meta-learners for every layer of the base-learner network 
        # (if layerwise is true, else we have one meta-learner network for the entire base-learner)
        self.metalearners = [MetaLearner(input_type=input_type, layers=layers, 
                                activation=act, history=ml_hist, time_input=time_input).to(self.dev) for layer in self.slices]


        # If input type is 'maml', the meta-learner network parametesr are fixed
        self.meta_params = list(self.initialization) if not self.decouple else []
        if not input_type == "maml":
            for ml in self.metalearners:
                self.meta_params += list(ml.parameters())
        
        # If we need to learn a learning rate per parameter
        if self.param_lr:
            self.lr_vector = torch.rand(size=(self.num_params, 1), device=self.dev).detach()
            self.lr_vector.requires_grad=True

            self.meta_params += [self.lr_vector]

        # Create gradient buffers that keep track of gradients across tasks
        self.grad_buffer = [torch.zeros(p.size(), device=self.dev) for p in self.meta_params]

        self.optimizer = self.opt_fn(self.meta_params, lr=self.lr)
        
        # All keys in the base learner model that have nothing to do with batch normalization
        self.keys = [k for k in self.baselearner.state_dict().keys()\
                     if not "running" in k and not "num" in k]
        
        self.bn_keys = [k for k in self.baselearner.state_dict().keys()\
                        if "running" in k or "num" in k]
        
    def _get_grad_buffers(self):
        """Initialize empty gradient buffers

        Create gradient buffers to store gradients across tasks.
        This allows for making outer-updates across batches of tasks.
        """

        # Create gradient buffers that keep track of gradients across tasks
        self.init_grad_buffer = [torch.zeros(p.size()) for p in self.initialization]
        self.meta_grad_buffers = []
        for meta_learner in self.metalearners:
            buffer = [torch.zeros(p.size()) for p in meta_learner.parameters()]
            self.meta_grad_buffers.append(buffer)
        
    
    def _history(self, H, step, grads, updates):
        """ Compute historical information tensor

        Obtain the appropriate historical information in tensor format that can be
        concatenated to the meta-learner input X. 

        Parameters
        ----------
        H : torch.Tensor
            Current historical information tensor
        step : int
            Step index (of inner-optimization procedure)
        grads : torch.Tensor
            Gradient tensor of loss w.r.t. all base-learner parameters
        updates : torch.Tensor
            Updates that were made in the previous step (initialized with 0 if no step has been made yet)
            
        Returns
        ----------
        torch.Tensor
            Updated tensor of historical information which can added to the meta-input tensor X
        """
        
        # If in the first step, we have to initialize the history
        if step == 0:
            # If our history type is gradients, initialize with the gradients
            if self.history_code == History.Gradients:
                H = grads
                return H
            # In case of updates, initialize with zeros because we didn't make any updates
            elif self.history_code == History.Updates:
                return torch.zeros(grads.size(), device=self.dev)

        # Initialize the updates with previously made update
        if step == 1 and self.history_code == History.Updates:
            return updates
        
        # Else, the step >= 1 H is no longer empty ==> simply update
        if self.history_code == History.Updates:
            new = updates
        else:
            new = grads

        new_H = self.beta * H + (1 - self.beta) * new
        return new_H

    def _deploy(self, train_x, train_y, test_x, test_y, return_network=False):
        """Run TURTLE on a single task to get the loss on the query set
        
        1. Evaluate the base-learner loss and gradients on the support set (train_x, train_y)
        using our initialization point.
        2. Make a single weight update based on this information.
        3. Evaluate and return the loss of the fast weights (initialization + proposed updates)
        
        Parameters
        ----------
        train_x : torch.Tensor
            Inputs of the support set
        train_y : torch.Tensor
            Outputs of the support set
        test_x : torch.Tensor
            Inputs of the query set
        test_y : torch.Tensor
            Outputs of the query set
            
        Returns
        ----------
        test_loss
            Loss of the base-learner on the query set after the proposed
            one-step update
        """

        H = None # Empty history
        updates = None # Empty updates
        fast_weights = [p.clone() for p in self.initialization]

        for step in range(self.T):
            # Construct meta input matrix
            # set_weights(self.baselearner, fast_weights, self.keys, self.bn_keys)
            X, grads = get_init_info(self.baselearner, train_x, train_y, weights=fast_weights,
                                 create_graph=self.second_order, 
                                 retain_graph=self.T > 1 or self.second_order,
                                 input_code=self.input_code, grad_clip=self.grad_clip)

            # If we use historical information
            if self.history_code != History.Empty:
                H = self._history(H, step, grads, updates)
                X = torch.cat((X, H), 1)

            # If time should be part of input, concatenate a column of <step> 
            # to X
            if self.time_input:
                col = torch.ones(len(X), 1, device=self.dev) * step
                X = torch.cat((X, col), 1)

            X_slices = [X[lb:ub] for (lb, ub) in self.slices]

            updates = None
            # Compute weight update proposals 
            # Slice the meta-input matrix X into pieces: one for every layer of the base-learner
            # (if layerwise is True, else we use one slice for the entire base-learner network)
            for lid, (lb, ub) in enumerate(self.slices):
                X_slice = X[lb:ub]
                if updates is None:
                    if not self.param_lr:
                        updates = self.metalearners[lid](X_slice)
                    else:
                        updates = torch.mul(self.metalearners[lid](X_slice), self.lr_vector[lb:ub])
                else:
                    if not self.param_lr:
                        updates = torch.cat([updates, self.metalearners[lid](X_slice)])
                    else:
                        updates = torch.cat([updates, torch.mul(self.metalearners[lid](X_slice), self.lr_vector[lb:ub])])
                    

            # Compute the task-specific weights by adding the updates
            # to our initialization
            fast_weights = get_fast_weights(updates=updates, 
                                           initialization=fast_weights)


        if return_network:
            return self.baselearner, fast_weights
        # Get and return performance on query set
        test_preds = self.baselearner.forward_weights(test_x, fast_weights)
        test_loss = self.baselearner.criterion(test_preds, test_y)
        return test_loss, test_preds


    def train(self, train_x, train_y, test_x, test_y):
        """Train on a given task
        
        1. Evaluate the base-learner loss and gradients on the support set (train_x, train_y)
        using our initialization point.
        2. Make a single weight update based on this information.
        3. Evaluate the loss of the fast weights (initialization + proposed updates)
        4. Backpropagete the loss to update the initialization weights and the OSOMetaLearner
           responsible for proposing the base-learner weight updates
        
        Parameters
        ----------
        train_x : torch.Tensor
            Inputs of the support set
        train_y : torch.Tensor
            Outputs of the support set
        test_x : torch.Tensor
            Inputs of the query set
        test_y : torch.Tensor
            Outputs of the query set
        """

        self.baselearner.train()
        self.task_counter += 1
        self.dcounter += 1
        # Put all tensors on right device
        # train_x, train_y, test_x, test_y = put_on_device(
        #                                     self.dev,
        #                                     [train_x, train_y,
        #                                     test_x, test_y])
        
        # Deploy TURTLE on task to compute test loss
        test_loss, _ = self._deploy(train_x, train_y, test_x, test_y)
       
        # Propagate loss backwards and update the initialization and
        # weights of the meta-learner
        test_loss.backward()

        # Clip gradients
        if not self.grad_clip is None:
            for p in self.initialization:
                p.grad = torch.clamp(p.grad, -self.grad_clip, +self.grad_clip)

        # self.init_grad_buffer = [self.init_grad_buffer[i] + self.initialization[i].grad for i in range(len(self.initialization))]
        # for i in range(len(self.meta_grad_buffers)):
        #     buffer = self.meta_grad_buffers[i]
        #     buffer = [buffer[j] + self.initialization[i].grad for i in range(len(self.initialization))]
        self.grad_buffer = [self.grad_buffer[i] + self.meta_params[i].grad for i in range(len(self.meta_params))]
        
        self.optimizer.zero_grad()
        for ml in self.metalearners:
            ml.zero_grad()

        if self.task_counter % self.meta_batch_size == 0:
            # Copy gradients from self.grad_buffer to gradient buffers in the initialization parameters
            for i, p in enumerate(self.meta_params):
                p.grad = self.grad_buffer[i]
            self.optimizer.step() 
            # Reset all stats for next update
            self.optimizer.zero_grad()
            for ml in self.metalearners:
                ml.zero_grad()
            self.grad_buffer = [torch.zeros(p.size(), device=self.dev) for p in self.meta_params]
            self.task_counter = 0

        if self.metasearch:
            if self.dcounter % 1000 == 0:
                print("Checkpoint -", self.dcounter)
            # pick new initialization with probability 0.1
            if np.random.rand() < 0.1:
                self.baselearner = self.baselearner_fn(**self.baselearner_args).to(self.dev)
                self.initialization, _= get_params(self.baselearner, self.dev)
                for p in self.initialization:
                    p.requires_grad = True
            
            # If switch limit has been reached, go over to initialization mode
            if self.dcounter == self.decouple:
                self.baselearner = self.baselearner_fn(**self.baselearner_args).to(self.dev)
                self.initialization, _= get_params(self.baselearner, self.dev)
                for p in self.initialization:
                    p.requires_grad = True
                
                self.meta_params = list(self.initialization)
                # Create gradient buffers that keep track of gradients across tasks
                self.grad_buffer = [torch.zeros(p.size()) for p in self.meta_params]
                self.optimizer = self.opt_fn(self.meta_params, lr=self.lr)

                self.metasearch = False
        
    def evaluate(self, train_x, train_y, test_x, test_y, return_network=False):
        """Evaluate on a given task
        
        Use the support set (train_x, train_y) to compute weight updates
        using the OSO. After the one-step update, evaluate and return the loss
        on the query set (test_x, test_y).
        
        Parameters
        ----------
        train_x : torch.Tensor
            Inputs of the support set
        train_y : torch.Tensor
            Outputs of the support set
        test_x : torch.Tensor
            Inputs of the query set
        test_y : torch.Tensor
            Outputs of the query set
            
        Returns
        ----------
        test_loss
            Loss/accuracy of the base-learner on the query set after the proposed
            one-step update
        """
        
        self.baselearner.eval()
        # Put all tensors on right device
        # train_x, train_y, test_x, test_y = put_on_device(
        #                                     self.dev,
        #                                     [train_x, train_y,
        #                                     test_x, test_y])

        # Deploy Turtle on task to compute test loss
        if return_network:
            return self._deploy(train_x, train_y, test_x, test_y, return_network)
        
        test_loss, preds = self._deploy(train_x, train_y, test_x, test_y, return_network)
            
        # If our objective is to minimize, return the loss,
        # else the accuracy
        if self.operator == min:
            return test_loss.item()
        else:
            # Turn one-hot predictions into class preds
            preds = torch.argmax(preds, dim=1)
            return accuracy(preds, test_y)
        
    def dump_state(self):
        """Return the state of the meta-learner
        
        Returns
        ----------
        initialization
            Initialization parameters
        state_dict meta-learner
            State dictionary of the meta-learner
        """
        
        # Copy initialization parameters
        init = [p.clone().detach() for p in self.initialization]
        return init, [ml.state_dict() for ml in self.metalearners]
    
    def load_state(self, state):
        """Load the given state into the meta-learner
        
        Parameters
        ----------
        state : [initialization, state_dict]
            Initialization parameters and state dictionary of the meta-learner
        """
        
        init, sd = state
        initialization = [p.clone() for p in init]
        for p in initialization:
            p.requires_grad = True
        self.initialization = initialization
        for i in range(len(sd)):
            self.metalearners[i].load_state_dict(sd[i])
<filename>baselines/turtle/turtle_code.py
import torch
import numpy as np

from algorithm import Algorithm
from shell import MetaLearner, get_init_info, get_fast_weights,\
                           INP_TO_CODE, HIST_TO_CODE, Input, History
from utils import put_on_device, get_init_score_and_operator, accuracy, get_params




class Turtle(Algorithm):
    """sTateless neURal meTa LEarning (TURTLE)
    
    Meta-learning algorithm that combines the best of MAML and 
    LSTM meta-learners. It learns both a good common initialization 
    point for tasks, and learns how to make parameter updates in a
    using a meta-learner network.
    
    ...

    Attributes
    ----------
    baselearner_fn : constructor function
        Constructor function for the base-learner
    baselearner_args : dict
        Dictionary of keyword arguments for the base-learner
    opt_fn : constructor function
        Constructor function for the optimizer to use
    T : int
        Number of update steps to parameters per task
    train_batch_size : int
        Indicating the size of minibatches that are sampled from meta-train tasks
    test_batch_size : int
        Size of batches to sample from meta-[val/test] tasks
    lr : float
        Learning rate for the optimizer
    validation : boolean
        Whether this model should use meta-validation
    dev : str
        Device identifier
    trainable : boolean
        Whether it can train on meta-trrain tasks
    episodic : boolean
        Whether to sample tasks or mini batches for training
        
    Methods
    -------
    train(train_x, train_y, test_x, test_y)
        Perform a single training step on a given task
    
    evaluate(train_x, train_y, test_x, test_y)
        Evaluate the performance on the given task
        
    dump_state()
        Dump the meta-learner state s.t. it can be loaded again later
        
    load_state(state)
        Set meta-learner state to provided @state 
    """
    
    def __init__(self, input_type, history, beta, grad_clip, layers, act, 
                 second_order, layer_wise, param_lr, decouple, meta_batch_size=1, time_input=False, **kwargs):
        """Set all (inherited) attributes
        
        Attributes
        ----------
        input_type : str
            Specifier of the input type (raw_grads, raw_loss_grads, proc_loss_grads,
            maml). Determines the input to the meta-learner network
        layers : list
            List of integers corresponding to the number of neurons per hidden/output layer
            e.g., [5,5,1] is a neural network [input -> hidden 5 -> hidden 5 -> output 1]
        history : str
            Type of historical information to use [none, grads, updates]
        beta : float
            Float between [0,1] that determines the memory capacity for the historical
            information. Only affects results if history != none
        grad_clip : float
            Gradient clipping value to use
        act : act_fn
            Activation function to user for meta-learner
        second_order : boolean
            Whether to use second-order gradient information
        layer_wise : boolean
            Whether to use a meta-learner network for every individual base-learner layer
        meta_batch_size : int
            Number of tasks to compute outer-update
        time_input : boolean
            Whether to add a timestamp to the input (this way TURTLE can find an equivalent of a learning schedule)
        param_lr : boolean
            Whether to maintain a learning rate per parameter
        decouple : int
            Whether to decouple the base- and meta-search processes (int indicates the task # to switch modes)
        **kwargs : dict
            Ignored keyword arguments
        """
        
        super().__init__(**kwargs)
        self.input_type = input_type
        self.input_code = INP_TO_CODE[input_type]
        self.history_code = HIST_TO_CODE[history]
        self.beta = beta
        self.grad_clip = grad_clip
        self.second_order = second_order
        self.layer_wise = layer_wise
        self.meta_batch_size = meta_batch_size
        self.time_input = time_input
        self.param_lr = param_lr
        self.decouple = decouple
        self.metasearch = bool(self.decouple) # currently looking only for optimizer if self.decouple
        self.dcounter = 0
        # Increment after every train step on a single task, and update
        # init when task_counter % meta_batch_size == 0
        self.task_counter = 0 

        # Get random initialization point for baselearner
        self.baselearner = self.baselearner_fn(**self.baselearner_args).to(self.dev)
        self.initialization, self.slices = get_params(self.baselearner, self.dev)

        # If we do not use layer-wise meta-learner networks, there is only a single parameter slice
        # namely (0, num_params)
        if not self.layer_wise:
            self.slices = [ tuple([0, self.slices[-1][1]]) ]
            print(self.slices) 

        self.num_params = 0
        # Enable gradient tracking for the initialization parameters    
        for p in self.initialization:
            self.num_params += p.numel()
            p.requires_grad = True

        # We do not use historical information when T = 1
        if self.T == 1:
            ml_hist = "none" 
        else:
            ml_hist = history

        # Create meta-learners for every layer of the base-learner network 
        # (if layerwise is true, else we have one meta-learner network for the entire base-learner)
        self.metalearners = [MetaLearner(input_type=input_type, layers=layers, 
                                activation=act, history=ml_hist, time_input=time_input).to(self.dev) for layer in self.slices]


        # If input type is 'maml', the meta-learner network parametesr are fixed
        self.meta_params = list(self.initialization) if not self.decouple else []
        if not input_type == "maml":
            for ml in self.metalearners:
                self.meta_params += list(ml.parameters())
        
        # If we need to learn a learning rate per parameter
        if self.param_lr:
            self.lr_vector = torch.rand(size=(self.num_params, 1), device=self.dev).detach()
            self.lr_vector.requires_grad=True

            self.meta_params += [self.lr_vector]

        # Create gradient buffers that keep track of gradients across tasks
        self.grad_buffer = [torch.zeros(p.size(), device=self.dev) for p in self.meta_params]

        self.optimizer = self.opt_fn(self.meta_params, lr=self.lr)
        
        # All keys in the base learner model that have nothing to do with batch normalization
        self.keys = [k for k in self.baselearner.state_dict().keys()\
                     if not "running" in k and not "num" in k]
        
        self.bn_keys = [k for k in self.baselearner.state_dict().keys()\
                        if "running" in k or "num" in k]
        
    def _get_grad_buffers(self):
        """Initialize empty gradient buffers

        Create gradient buffers to store gradients across tasks.
        This allows for making outer-updates across batches of tasks.
        """

        # Create gradient buffers that keep track of gradients across tasks
        self.init_grad_buffer = [torch.zeros(p.size()) for p in self.initialization]
        self.meta_grad_buffers = []
        for meta_learner in self.metalearners:
            buffer = [torch.zeros(p.size()) for p in meta_learner.parameters()]
            self.meta_grad_buffers.append(buffer)
        
    
    def _history(self, H, step, grads, updates):
        """ Compute historical information tensor

        Obtain the appropriate historical information in tensor format that can be
        concatenated to the meta-learner input X. 

        Parameters
        ----------
        H : torch.Tensor
            Current historical information tensor
        step : int
            Step index (of inner-optimization procedure)
        grads : torch.Tensor
            Gradient tensor of loss w.r.t. all base-learner parameters
        updates : torch.Tensor
            Updates that were made in the previous step (initialized with 0 if no step has been made yet)
            
        Returns
        ----------
        torch.Tensor
            Updated tensor of historical information which can added to the meta-input tensor X
        """
        
        # If in the first step, we have to initialize the history
        if step == 0:
            # If our history type is gradients, initialize with the gradients
            if self.history_code == History.Gradients:
                H = grads
                return H
            # In case of updates, initialize with zeros because we didn't make any updates
            elif self.history_code == History.Updates:
                return torch.zeros(grads.size(), device=self.dev)

        # Initialize the updates with previously made update
        if step == 1 and self.history_code == History.Updates:
            return updates
        
        # Else, the step >= 1 H is no longer empty ==> simply update
        if self.history_code == History.Updates:
            new = updates
        else:
            new = grads

        new_H = self.beta * H + (1 - self.beta) * new
        return new_H

    def _deploy(self, train_x, train_y, test_x, test_y, return_network=False):
        """Run TURTLE on a single task to get the loss on the query set
        
        1. Evaluate the base-learner loss and gradients on the support set (train_x, train_y)
        using our initialization point.
        2. Make a single weight update based on this information.
        3. Evaluate and return the loss of the fast weights (initialization + proposed updates)
        
        Parameters
        ----------
        train_x : torch.Tensor
            Inputs of the support set
        train_y : torch.Tensor
            Outputs of the support set
        test_x : torch.Tensor
            Inputs of the query set
        test_y : torch.Tensor
            Outputs of the query set
            
        Returns
        ----------
        test_loss
            Loss of the base-learner on the query set after the proposed
            one-step update
        """

        H = None # Empty history
        updates = None # Empty updates
        fast_weights = [p.clone() for p in self.initialization]

        for step in range(self.T):
            # Construct meta input matrix
            # set_weights(self.baselearner, fast_weights, self.keys, self.bn_keys)
            X, grads = get_init_info(self.baselearner, train_x, train_y, weights=fast_weights,
                                 create_graph=self.second_order, 
                                 retain_graph=self.T > 1 or self.second_order,
                                 input_code=self.input_code, grad_clip=self.grad_clip)

            # If we use historical information
            if self.history_code != History.Empty:
                H = self._history(H, step, grads, updates)
                X = torch.cat((X, H), 1)

            # If time should be part of input, concatenate a column of <step> 
            # to X
            if self.time_input:
                col = torch.ones(len(X), 1, device=self.dev) * step
                X = torch.cat((X, col), 1)

            X_slices = [X[lb:ub] for (lb, ub) in self.slices]

            updates = None
            # Compute weight update proposals 
            # Slice the meta-input matrix X into pieces: one for every layer of the base-learner
            # (if layerwise is True, else we use one slice for the entire base-learner network)
            for lid, (lb, ub) in enumerate(self.slices):
                X_slice = X[lb:ub]
                if updates is None:
                    if not self.param_lr:
                        updates = self.metalearners[lid](X_slice)
                    else:
                        updates = torch.mul(self.metalearners[lid](X_slice), self.lr_vector[lb:ub])
                else:
                    if not self.param_lr:
                        updates = torch.cat([updates, self.metalearners[lid](X_slice)])
                    else:
                        updates = torch.cat([updates, torch.mul(self.metalearners[lid](X_slice), self.lr_vector[lb:ub])])
                    

            # Compute the task-specific weights by adding the updates
            # to our initialization
            fast_weights = get_fast_weights(updates=updates, 
                                           initialization=fast_weights)


        if return_network:
            return self.baselearner, fast_weights
        # Get and return performance on query set
        test_preds = self.baselearner.forward_weights(test_x, fast_weights)
        test_loss = self.baselearner.criterion(test_preds, test_y)
        return test_loss, test_preds


    def train(self, train_x, train_y, test_x, test_y):
        """Train on a given task
        
        1. Evaluate the base-learner loss and gradients on the support set (train_x, train_y)
        using our initialization point.
        2. Make a single weight update based on this information.
        3. Evaluate the loss of the fast weights (initialization + proposed updates)
        4. Backpropagete the loss to update the initialization weights and the OSOMetaLearner
           responsible for proposing the base-learner weight updates
        
        Parameters
        ----------
        train_x : torch.Tensor
            Inputs of the support set
        train_y : torch.Tensor
            Outputs of the support set
        test_x : torch.Tensor
            Inputs of the query set
        test_y : torch.Tensor
            Outputs of the query set
        """

        self.baselearner.train()
        self.task_counter += 1
        self.dcounter += 1
        # Put all tensors on right device
        # train_x, train_y, test_x, test_y = put_on_device(
        #                                     self.dev,
        #                                     [train_x, train_y,
        #                                     test_x, test_y])
        
        # Deploy TURTLE on task to compute test loss
        test_loss, _ = self._deploy(train_x, train_y, test_x, test_y)
       
        # Propagate loss backwards and update the initialization and
        # weights of the meta-learner
        test_loss.backward()

        # Clip gradients
        if not self.grad_clip is None:
            for p in self.initialization:
                p.grad = torch.clamp(p.grad, -self.grad_clip, +self.grad_clip)

        # self.init_grad_buffer = [self.init_grad_buffer[i] + self.initialization[i].grad for i in range(len(self.initialization))]
        # for i in range(len(self.meta_grad_buffers)):
        #     buffer = self.meta_grad_buffers[i]
        #     buffer = [buffer[j] + self.initialization[i].grad for i in range(len(self.initialization))]
        self.grad_buffer = [self.grad_buffer[i] + self.meta_params[i].grad for i in range(len(self.meta_params))]
        
        self.optimizer.zero_grad()
        for ml in self.metalearners:
            ml.zero_grad()

        if self.task_counter % self.meta_batch_size == 0:
            # Copy gradients from self.grad_buffer to gradient buffers in the initialization parameters
            for i, p in enumerate(self.meta_params):
                p.grad = self.grad_buffer[i]
            self.optimizer.step() 
            # Reset all stats for next update
            self.optimizer.zero_grad()
            for ml in self.metalearners:
                ml.zero_grad()
            self.grad_buffer = [torch.zeros(p.size(), device=self.dev) for p in self.meta_params]
            self.task_counter = 0

        if self.metasearch:
            if self.dcounter % 1000 == 0:
                print("Checkpoint -", self.dcounter)
            # pick new initialization with probability 0.1
            if np.random.rand() < 0.1:
                self.baselearner = self.baselearner_fn(**self.baselearner_args).to(self.dev)
                self.initialization, _= get_params(self.baselearner, self.dev)
                for p in self.initialization:
                    p.requires_grad = True
            
            # If switch limit has been reached, go over to initialization mode
            if self.dcounter == self.decouple:
                self.baselearner = self.baselearner_fn(**self.baselearner_args).to(self.dev)
                self.initialization, _= get_params(self.baselearner, self.dev)
                for p in self.initialization:
                    p.requires_grad = True
                
                self.meta_params = list(self.initialization)
                # Create gradient buffers that keep track of gradients across tasks
                self.grad_buffer = [torch.zeros(p.size()) for p in self.meta_params]
                self.optimizer = self.opt_fn(self.meta_params, lr=self.lr)

                self.metasearch = False
        
    def evaluate(self, train_x, train_y, test_x, test_y, return_network=False):
        """Evaluate on a given task
        
        Use the support set (train_x, train_y) to compute weight updates
        using the OSO. After the one-step update, evaluate and return the loss
        on the query set (test_x, test_y).
        
        Parameters
        ----------
        train_x : torch.Tensor
            Inputs of the support set
        train_y : torch.Tensor
            Outputs of the support set
        test_x : torch.Tensor
            Inputs of the query set
        test_y : torch.Tensor
            Outputs of the query set
            
        Returns
        ----------
        test_loss
            Loss/accuracy of the base-learner on the query set after the proposed
            one-step update
        """
        
        self.baselearner.eval()
        # Put all tensors on right device
        # train_x, train_y, test_x, test_y = put_on_device(
        #                                     self.dev,
        #                                     [train_x, train_y,
        #                                     test_x, test_y])

        # Deploy Turtle on task to compute test loss
        if return_network:
            return self._deploy(train_x, train_y, test_x, test_y, return_network)
        
        test_loss, preds = self._deploy(train_x, train_y, test_x, test_y, return_network)
            
        # If our objective is to minimize, return the loss,
        # else the accuracy
        if self.operator == min:
            return test_loss.item()
        else:
            # Turn one-hot predictions into class preds
            preds = torch.argmax(preds, dim=1)
            return accuracy(preds, test_y)
        
    def dump_state(self):
        """Return the state of the meta-learner
        
        Returns
        ----------
        initialization
            Initialization parameters
        state_dict meta-learner
            State dictionary of the meta-learner
        """
        
        # Copy initialization parameters
        init = [p.clone().detach() for p in self.initialization]
        return init, [ml.state_dict() for ml in self.metalearners]
    
    def load_state(self, state):
        """Load the given state into the meta-learner
        
        Parameters
        ----------
        state : [initialization, state_dict]
            Initialization parameters and state dictionary of the meta-learner
        """
        
        init, sd = state
        initialization = [p.clone() for p in init]
        for p in initialization:
            p.requires_grad = True
        self.initialization = initialization
        for i in range(len(sd)):
            self.metalearners[i].load_state_dict(sd[i])
sTateless neURal meTa LEarning (TURTLE)          Meta-learning algorithm that combines the best of MAML and      LSTM meta-learners. It learns both a good common initialization      point for tasks, and learns how to make parameter updates in a     using a meta-learner network.          ...      Attributes     ----------     baselearner_fn : constructor function         Constructor function for the base-learner     baselearner_args : dict         Dictionary of keyword arguments for the base-learner     opt_fn : constructor function         Constructor function for the optimizer to use     T : int         Number of update steps to parameters per task     train_batch_size : int         Indicating the size of minibatches that are sampled from meta-train tasks     test_batch_size : int         Size of batches to sample from meta-[val/test] tasks     lr : float         Learning rate for the optimizer     validation : boolean         Whether this model should use meta-validation     dev : str         Device identifier     trainable : boolean         Whether it can train on meta-trrain tasks     episodic : boolean         Whether to sample tasks or mini batches for training              Methods     -------     train(train_x, train_y, test_x, test_y)         Perform a single training step on a given task          evaluate(train_x, train_y, test_x, test_y)         Evaluate the performance on the given task              dump_state()         Dump the meta-learner state s.t. it can be loaded again later              load_state(state)         Set meta-learner state to provided @state Set all (inherited) attributes                  Attributes         ----------         input_type : str             Specifier of the input type (raw_grads, raw_loss_grads, proc_loss_grads,             maml). Determines the input to the meta-learner network         layers : list             List of integers corresponding to the number of neurons per hidden/output layer             e.g., [5,5,1] is a neural network [input -> hidden 5 -> hidden 5 -> output 1]         history : str             Type of historical information to use [none, grads, updates]         beta : float             Float between [0,1] that determines the memory capacity for the historical             information. Only affects results if history != none         grad_clip : float             Gradient clipping value to use         act : act_fn             Activation function to user for meta-learner         second_order : boolean             Whether to use second-order gradient information         layer_wise : boolean             Whether to use a meta-learner network for every individual base-learner layer         meta_batch_size : int             Number of tasks to compute outer-update         time_input : boolean             Whether to add a timestamp to the input (this way TURTLE can find an equivalent of a learning schedule)         param_lr : boolean             Whether to maintain a learning rate per parameter         decouple : int             Whether to decouple the base- and meta-search processes (int indicates the task # to switch modes)         **kwargs : dict             Ignored keyword arguments # currently looking only for optimizer if self.decouple # Increment after every train step on a single task, and update # init when task_counter % meta_batch_size == 0 # Get random initialization point for baselearner # If we do not use layer-wise meta-learner networks, there is only a single parameter slice # namely (0, num_params) # Enable gradient tracking for the initialization parameters # We do not use historical information when T = 1 # Create meta-learners for every layer of the base-learner network # (if layerwise is true, else we have one meta-learner network for the entire base-learner) # If input type is 'maml', the meta-learner network parametesr are fixed # If we need to learn a learning rate per parameter # Create gradient buffers that keep track of gradients across tasks # All keys in the base learner model that have nothing to do with batch normalization Initialize empty gradient buffers          Create gradient buffers to store gradients across tasks.         This allows for making outer-updates across batches of tasks. # Create gradient buffers that keep track of gradients across tasks Compute historical information tensor          Obtain the appropriate historical information in tensor format that can be         concatenated to the meta-learner input X.           Parameters         ----------         H : torch.Tensor             Current historical information tensor         step : int             Step index (of inner-optimization procedure)         grads : torch.Tensor             Gradient tensor of loss w.r.t. all base-learner parameters         updates : torch.Tensor             Updates that were made in the previous step (initialized with 0 if no step has been made yet)                      Returns         ----------         torch.Tensor             Updated tensor of historical information which can added to the meta-input tensor X # If in the first step, we have to initialize the history # If our history type is gradients, initialize with the gradients # In case of updates, initialize with zeros because we didn't make any updates # Initialize the updates with previously made update # Else, the step >= 1 H is no longer empty ==> simply update Run TURTLE on a single task to get the loss on the query set                  1. Evaluate the base-learner loss and gradients on the support set (train_x, train_y)         using our initialization point.         2. Make a single weight update based on this information.         3. Evaluate and return the loss of the fast weights (initialization + proposed updates)                  Parameters         ----------         train_x : torch.Tensor             Inputs of the support set         train_y : torch.Tensor             Outputs of the support set         test_x : torch.Tensor             Inputs of the query set         test_y : torch.Tensor             Outputs of the query set                      Returns         ----------         test_loss             Loss of the base-learner on the query set after the proposed             one-step update # Empty history # Empty updates # Construct meta input matrix # set_weights(self.baselearner, fast_weights, self.keys, self.bn_keys) # If we use historical information # If time should be part of input, concatenate a column of <step> # to X # Compute weight update proposals # Slice the meta-input matrix X into pieces: one for every layer of the base-learner # (if layerwise is True, else we use one slice for the entire base-learner network) # Compute the task-specific weights by adding the updates # to our initialization # Get and return performance on query set Train on a given task                  1. Evaluate the base-learner loss and gradients on the support set (train_x, train_y)         using our initialization point.         2. Make a single weight update based on this information.         3. Evaluate the loss of the fast weights (initialization + proposed updates)         4. Backpropagete the loss to update the initialization weights and the OSOMetaLearner            responsible for proposing the base-learner weight updates                  Parameters         ----------         train_x : torch.Tensor             Inputs of the support set         train_y : torch.Tensor             Outputs of the support set         test_x : torch.Tensor             Inputs of the query set         test_y : torch.Tensor             Outputs of the query set # Put all tensors on right device # train_x, train_y, test_x, test_y = put_on_device( #                                     self.dev, #                                     [train_x, train_y, #                                     test_x, test_y]) # Deploy TURTLE on task to compute test loss # Propagate loss backwards and update the initialization and # weights of the meta-learner # Clip gradients # self.init_grad_buffer = [self.init_grad_buffer[i] + self.initialization[i].grad for i in range(len(self.initialization))] # for i in range(len(self.meta_grad_buffers)): #     buffer = self.meta_grad_buffers[i] #     buffer = [buffer[j] + self.initialization[i].grad for i in range(len(self.initialization))] # Copy gradients from self.grad_buffer to gradient buffers in the initialization parameters # Reset all stats for next update # pick new initialization with probability 0.1 # If switch limit has been reached, go over to initialization mode # Create gradient buffers that keep track of gradients across tasks Evaluate on a given task                  Use the support set (train_x, train_y) to compute weight updates         using the OSO. After the one-step update, evaluate and return the loss         on the query set (test_x, test_y).                  Parameters         ----------         train_x : torch.Tensor             Inputs of the support set         train_y : torch.Tensor             Outputs of the support set         test_x : torch.Tensor             Inputs of the query set         test_y : torch.Tensor             Outputs of the query set                      Returns         ----------         test_loss             Loss/accuracy of the base-learner on the query set after the proposed             one-step update # Put all tensors on right device # train_x, train_y, test_x, test_y = put_on_device( #                                     self.dev, #                                     [train_x, train_y, #                                     test_x, test_y]) # Deploy Turtle on task to compute test loss # If our objective is to minimize, return the loss, # else the accuracy # Turn one-hot predictions into class preds Return the state of the meta-learner                  Returns         ----------         initialization             Initialization parameters         state_dict meta-learner             State dictionary of the meta-learner # Copy initialization parameters Load the given state into the meta-learner                  Parameters         ----------         state : [initialization, state_dict]             Initialization parameters and state dictionary of the meta-learner