我的模型代码如下，我想在其中加入inner_loop_optimizers模块，使其能够自学习地调节学习率和模型权重
import time
import os
import sys
import argparse
import shutil
import numpy as np
import torch.utils
from tqdm import tqdm
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torch.optim import lr_scheduler
from tensorboardX import SummaryWriter
from easydict import EasyDict
from network import model_selection_meta as model_selection
from torch.utils.data.sampler import SubsetRandomSampler, RandomSampler
import random
# from network import FNet
from utils.loss import CompactLoss
from torch.utils.data.dataset import Dataset
from torch.nn import functional as F
import copy
from utils.config import *
from utils.roc import cal_metric
from utils import *
import pdb
from PIL import Image

from inner_loop_optimizers import LSLRGradientDescentLearningRule
import higher

def worker_init_fn(x):
    seed = RNG_SEED
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed)
    return

def model_forward(image, model, post_function=nn.Sigmoid(),feat = False, params=None):
    """
    Predicts the label of an input image. Preprocesses the input image and
    casts it to cuda if required

    :param image: numpy image
    :param model: torch model with linear layer at the end
    :param post_function: e.g., softmax
    :return: prediction (1 = fake, 0 = real)
    """

    # Model prediction
    output,feature = model(image, params=params)
    output = post_function(output)
    output = output.squeeze()
    
    prediction = (output >= 0.5).float()
    if feat==False:
        return prediction, output
    else:
        return prediction, output,feature

class RegularizerNetwork(nn.Module):
    def __init__(self, embedding_size, num_layers):
        super(RegularizerNetwork, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(embedding_size, embedding_size),
            nn.ReLU(inplace=True),
            nn.Linear(embedding_size, num_layers * 2)  # 输出 alpha 和 beta
        )
    
    def forward(self, embeddings):
        output = self.fc(embeddings)
        num_layers = output.size(-1) // 2
        alphas = output[:, :num_layers]
        betas = output[:, num_layers:]
        return alphas, betas



def train(model, optimizer, train_dataloader, meta_dataloader, criterion_oc, epoch, epoch_size, device, inner_loop_optimizer, regularizer_network):
    model.train()
    regularizer_network.train()
    inner_loop_optimizer.train()

    if meta_dataloader is not None:
        meta_dataloader_iter = iter(meta_dataloader)
    losses = AverageMeter()
    acces = AverageMeter()
    meta_losses = AverageMeter()
    meta_acces = AverageMeter()
    post_function = nn.Sigmoid()
    criterion = torch.nn.BCELoss(reduction='mean').to(device)
    num_inner_steps = inner_loop_optimizer.total_num_inner_loop_steps
    print_interval = 10  # Adjust as needed

    for i, datas in enumerate(train_dataloader):
        images = datas[0].to(device)
        targets = datas[1].float().to(device)

        optimizer.zero_grad()

        # Use higher to create a differentiable version of the model
        with higher.innerloop_ctx(model, inner_loop_optimizer, copy_initial_weights=False) as (fmodel, diffopt):
            # Inner loop: multiple steps
            for inner_step in range(num_inner_steps):
                # Forward pass
                prediction, output, feature = model_forward(images, fmodel, feat=True)
                compact_loss = criterion_oc(output, targets)
                cost = criterion(output, targets)
                support_loss = cost + lamda * compact_loss

                # Compute task embedding
                task_embedding = feature.mean(dim=0, keepdim=True)

                # Generate alpha and beta
                alphas, betas = regularizer_network(task_embedding)
                alphas = alphas.squeeze()
                betas = betas.squeeze()

                # Update inner loop optimizer parameters (if needed)
                # Assume inner_loop_optimizer has methods to set alphas and betas
                inner_loop_optimizer.set_alphas_betas(alphas, betas)

                # Update parameters using the differentiable optimizer
                diffopt.step(support_loss)

            # Meta-update using the updated model on meta-data
            if meta_dataloader is not None:
                try:
                    meta_datas = next(meta_dataloader_iter)
                except StopIteration:
                    meta_dataloader_iter = iter(meta_dataloader)
                    meta_datas = next(meta_dataloader_iter)
                meta_images = meta_datas[0].to(device)
                meta_targets = meta_datas[1].float().to(device)
            else:
                meta_images = images
                meta_targets = targets

            # Forward pass with the updated model
            prediction_meta, output_meta = model_forward(meta_images, fmodel)
            meta_loss = criterion(output_meta, meta_targets)

        # Backward pass and update the outer loop parameters
        meta_loss.backward()
        optimizer.step()

        # Compute and record metrics
        acc = (prediction.detach() >= 0.5).float().eq(targets).float().mean()
        loss = support_loss.item()
        meta_acc = (prediction_meta.detach() >= 0.5).float().eq(meta_targets).float().mean()

        meta_acces.update(meta_acc.item(), meta_targets.size(0))
        acces.update(acc.item(), targets.size(0))
        losses.update(loss, targets.size(0))
        meta_losses.update(meta_loss.item(), meta_targets.size(0))

        if i % print_interval == 0:
            print(f'{time.ctime()} || Epoch:{epoch} || Iter:{i}/{epoch_size} || '
                  f'Loss:{losses.val:.4f}({losses.avg:.4f}) || '
                  f'Accuracy:{acces.val:.4f}({acces.avg:.4f}) || '
                  f'MetaAccuracy:{meta_acces.val:.4f}({meta_acces.avg:.4f}) || '
                  f'MetaLoss:{meta_losses.val:.4f}({meta_losses.avg:.4f})')

def test(data_loader, model, device):
    model.eval()
    acces = []
    losses = 0.0
    label_list = []
    output_list = []
    criterion = torch.nn.BCELoss().to(device)
    wrongimg = []
    for i, datas in enumerate(tqdm(data_loader)):
        images = datas[0].to(device)#3,3,224,224
        targets = datas[1].float().to(device)
        
        with torch.no_grad():
            prediction, output = model_forward(images, model)
        label_list.extend(targets.cpu().numpy().tolist())
        output_list.extend(output.cpu().numpy().tolist())

        # acces.append((targets == prediction).cpu().numpy())
        batch_accuracy = (targets == prediction).float().mean().cpu().numpy()
        acces.append(batch_accuracy)
        loss = criterion(output, targets).item()
        losses += loss
    metrics = EasyDict()
    metrics.acc = np.mean(acces)
    eer,TPRs, auc,scaler = cal_metric(label_list,output_list,False)

    metrics.loss = losses / len(data_loader)
    metrics.auc = auc
    metrics.tpr = TPRs
    metrics.eer = eer
    model.train()
    
    return metrics

def updatebest(metrics,best_acc,best_loss,name,model,pnet=None):
    if metrics.acc>best_acc:
        best_acc = metrics.acc
        if model!=None:
            save_checkpoint(model.state_dict(), fpath=f'{save_dir}/{model_name}_best{name}acc.pth')
        if pnet!=None:
            save_checkpoint(pnet.state_dict(), fpath=f'{save_dir}/{model_name}_pnet_best{name}acc.pth')
        print(f'best_{name}_acc:{best_acc:.5f} (updated)', end=' ')
    else:
        print(f'best_{name}_acc:{best_acc:.5f} ')

    if metrics.loss<best_loss:
        best_loss = metrics.loss
        if model!=None:
            save_checkpoint(model.state_dict(), fpath=f'{save_dir}/{model_name}_best{name}loss.pth')
        if pnet!=None:
            save_checkpoint(pnet.state_dict(), fpath=f'{save_dir}/{model_name}_pnet_best{name}loss.pth')
        print(f'best_{name}_loss:{best_loss:.5f} (updated)')
    else:
        print(f'best_{name}_loss:{best_loss:.5f} ')

    return best_acc,best_loss

class FFDIDataset(Dataset):
    def __init__(self, img_path, img_label, transform=None):
        self.img_path = img_path
        self.img_label = img_label
        
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None
    
    def __getitem__(self, index):
        img = Image.open(self.img_path[index]).convert('RGB')
        
        if self.transform is not None:
            img = self.transform(img)
        
        return img, torch.from_numpy(np.array(self.img_label[index]))
    
    def __len__(self):
        return len(self.img_path)

def main():
    if os.path.exists(save_dir):
        shutil.rmtree(save_dir)
    print(f'save dir :{save_dir}')
    sys.stdout = Logger(os.path.join(save_dir, 'train.log'))

    device = 'cuda' if torch.cuda.is_available else 'cpu'

    model = model_selection(model_name=model_name, num_classes=1)

    # fnet = FNet(model.num_ftrs).to(device)


    if model_path is not None:    
        state_dict = torch.load(model_path)
        model.load_state_dict(state_dict)
        print(f'resume model from {model_path}')
    else:
        print('No model found, initializing random model.')
    if fnet_path is not None:    
        state_dict = torch.load(fnet_path)
        # fnet.load_state_dict(state_dict)
        print(f'resume model from {fnet_path}')
    else:
        print('No fnet_model found, initializing random model.')

    
    model = model.to(device)
    if parallel:
        model = nn.DataParallel(model)
    # 初始化正则化器网络
    embedding_size = model.module.num_ftrs if parallel else model.num_ftrs
    num_layers = len(list(model.named_parameters()))
    # 初始化内循环优化器
    init_inner_loop_learning_rate = 1e-3  
    init_weight_decay = 1e-4  
    num_inner_steps = 5
    inner_loop_optimizer = LSLRGradientDescentLearningRule(
        device=device,
        total_num_inner_loop_steps=num_inner_steps,
        use_learnable_learning_rates=True,
        use_learnable_weight_decay=True,
        alfa=True,
        random_init=False,
        init_learning_rate=init_inner_loop_learning_rate,
        init_weight_decay=init_weight_decay
    )
    regularizer_network = RegularizerNetwork(embedding_size, num_layers).to(device)
    # 将正则化器网络和内循环优化器的参数添加到优化器中
    optimizer = torch.optim.Adam(
        list(model.parameters()) + list(regularizer_network.parameters()) + list(inner_loop_optimizer.parameters()),
        lr=learning_rate, betas=(beta1, beta2)
    )

    scheduler = lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)

    criterion = torch.nn.BCELoss().to(device)
    criterion_oc = CompactLoss().to(device)



    train_label = pd.read_csv('/root/autodl-tmp/FaceForensics++/c23/train.txt', sep=' ', header=None, names=['img_name', 'target'])
    val_label = pd.read_csv('/root/autodl-tmp/FaceForensics++/c23/val.txt', sep=' ', header=None, names=['img_name', 'target'])
    train_label['path'] = train_label['img_name']
    val_label['path'] = val_label['img_name']
    
    support_label = pd.read_csv('/root/autodl-tmp/Celeb-DF-v2/test.txt', sep=' ', header=None, names=['img_name', 'target'])
    support_label['path'] = support_label['img_name']
    FFval_dataset = FFDIDataset(val_label['path'], val_label['target'], 
            transforms.Compose([
                        transforms.Resize((224, 224)),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                        ])
            )

    FF_test_dataloader = torch.utils.data.DataLoader(FFval_dataset, batch_size=batch_size, shuffle=False, num_workers=4)


    model.train()
    
    best_acc = 0.
    best_loss = 1000.

    for epoch in range(epochs):
        '''
        random shuffile all the source domain and split it into training domain and meta domain
        '''
        # copydatalist = copy.deepcopy(datasetlist)
        # random.seed(epoch)
        # random.shuffle(copydatalist)
        # meta_dataset = copydatalist[2].cat(copydatalist[2],randomseed = epoch)

        # copydatalist[0].set_meta_type(copydatalist[2].type)
        # copydatalist[1].set_meta_type(copydatalist[2].type)

        # train_dataset = copydatalist[0].cat(copydatalist[1])
        sampled_data = train_label.sample(n=10000).reset_index(drop=True)
        support_data = support_label.sample(n=16).reset_index(drop=True)
        # train_dataset = FFDIDataset(train_label['path'], train_label['target'],     
        train_dataset = FFDIDataset(sampled_data['path'], sampled_data['target'],        
            transforms.Compose([
                        transforms.Resize((224, 224)),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                        ])
        )
        support_dataset = FFDIDataset(support_data['path'], support_label['target'], 
            transforms.Compose([
                        transforms.Resize((224, 224)),
                        transforms.ToTensor(),
                        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                        ])
            )
        epoch_size = len(train_dataset) //batch_size
        # print(f"train dataset is:{copydatalist[0].type},{copydatalist[1].type},meta dataset is:{meta_dataset.type}")
        train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4,worker_init_fn=worker_init_fn)
        support_dataloader = torch.utils.data.DataLoader(support_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
        train(model, optimizer, train_dataloader, support_dataloader, criterion_oc,epoch, epoch_size,device, inner_loop_optimizer, regularizer_network)

        scheduler.step()
        # scheduler_fnet.step()

        if (epoch + 1) % test_interval == 0:
            '''
            Testing model on serval test-set.
            '''
            FF_metrics = test(FF_test_dataloader, model, device)
            metrics_list = [FF_metrics]
            avg_metrics = EasyDict()
            all_avg_metrics = EasyDict()
            avg_metrics.acc = (FF_metrics.acc)
            avg_metrics.loss = (FF_metrics.loss)
            avg_metrics.auc = (FF_metrics.auc)
            avg_metrics.eer = (FF_metrics.eer)

            print(f"avg acc:{avg_metrics.acc:.3f},loss:{avg_metrics.loss:.3f},auc:{avg_metrics.auc:.3f},eer:{avg_metrics.eer:.3f}")
            best_acc,best_loss = updatebest(avg_metrics,best_acc,best_loss,"avg",model)

            save_checkpoint(model.state_dict(), fpath=f'{save_dir}/{model_name}_lastepoch.pth')
            # save_checkpoint(fnet.state_dict(), fpath=f'{save_dir}/{model_name}_pnet_lastepoch.pth')
            
    print(f'save dir :{save_dir} done!!!')

if __name__ == '__main__':
    if DETERMINSTIC:
        np.random.seed(RNG_SEED)         
        torch.manual_seed(RNG_SEED)
        torch.cuda.manual_seed(RNG_SEED)           
        torch.cuda.manual_seed_all(RNG_SEED)         
        random.seed(RNG_SEED)         
        torch.backends.cudnn.benchmark = False         
        torch.backends.cudnn.deterministic = True         
        torch.backends.cudnn.enabled = True   
    main()
其中我的inner_loop_optimizers模块如下面所示import logging
import os
from collections import OrderedDict

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class LSLRGradientDescentLearningRule(nn.Module):
    """Simple (stochastic) gradient descent learning rule.
    For a scalar error function `E(p[0], p_[1] ... )` of some set of
    potentially multidimensional parameters this attempts to find a local
    minimum of the loss function by applying updates to each parameter of the
    form
        p[i] := p[i] - learning_rate * dE/dp[i]
    With `learning_rate` a positive scaling parameter.
    The error function used in successive applications of these updates may be
    a stochastic estimator of the true error function (e.g. when the error with
    respect to only a subset of data-points is calculated) in which case this
    will correspond to a stochastic gradient descent learning rule.
    """

    def __init__(self, device, total_num_inner_loop_steps, use_learnable_learning_rates, use_learnable_weight_decay, alfa, random_init, init_learning_rate=1e-3, init_weight_decay=5e-4):
        """Creates a new learning rule object.
        Args:
            init_learning_rate: A postive scalar to scale gradient updates to the
                parameters by. This needs to be carefully set - if too large
                the learning dynamic will be unstable and may diverge, while
                if set too small learning will proceed very slowly.
        """
        super(LSLRGradientDescentLearningRule, self).__init__()
        print(init_learning_rate)
        assert init_learning_rate > 0., 'learning_rate should be positive.'

        self.alfa = alfa
        self.random_init = random_init

        self.init_lr_val = init_learning_rate
        self.init_wd_val = init_weight_decay

        self.init_pspl_weight_decay = torch.ones(1)
        self.init_pspl_weight_decay.to(device)

        self.init_learning_rate = torch.ones(1) * init_learning_rate
        self.init_learning_rate.to(device)
        self.total_num_inner_loop_steps = total_num_inner_loop_steps
        self.use_learnable_learning_rates = use_learnable_learning_rates
        self.use_learnable_weight_decay = use_learnable_weight_decay
        self.init_weight_decay = torch.ones(1) * init_weight_decay
        self.init_bias_decay = torch.ones(1) 

    def initialise(self, names_weights_dict):
        if self.alfa:
            if self.random_init:
                self.names_beta_dict_per_param = nn.ParameterDict()

            self.names_alpha_dict = nn.ParameterDict()
            self.names_beta_dict = nn.ParameterDict()

            for idx, (key, param) in enumerate(names_weights_dict.items()):

                if self.random_init:
                # per-param weight decay for random init
                    self.names_beta_dict_per_param[key.replace(".", "-")] = nn.Parameter(
                        data=torch.ones(param.shape) * self.init_weight_decay * self.init_learning_rate,
                        requires_grad=self.use_learnable_learning_rates)

                    self.names_beta_dict[key.replace(".", "-")] = nn.Parameter(
                        data=torch.ones(self.total_num_inner_loop_steps + 1),
                        requires_grad=self.use_learnable_learning_rates)
                else:
                    # per-step per-layer meta-learnable weight decay bias term (for more stable training and better performance by 2~3%)
                    self.names_beta_dict[key.replace(".", "-")] = nn.Parameter(
                        data=torch.ones(self.total_num_inner_loop_steps + 1) * self.init_weight_decay * self.init_learning_rate,
                        requires_grad=self.use_learnable_learning_rates)

                # per-step per-layer meta-learnable learning rate bias term (for more stable training and better performance by 2~3%)
                self.names_alpha_dict[key.replace(".", "-")] = nn.Parameter(
                    data=torch.ones(self.total_num_inner_loop_steps + 1) * self.init_learning_rate,
                    requires_grad=self.use_learnable_learning_rates)
 
    def update_params(self, names_weights_dict, names_grads_wrt_params_dict, generated_alpha_params, generated_beta_params, num_step, tau=0.1):
        """Applies a single gradient descent update to all parameters.
        All parameter updates are performed using in-place operations and so
        nothing is returned.
        Args:
            grads_wrt_params: A list of gradients of the scalar loss function
                with respect to each of the parameters passed to `initialise`
                previously, with this list expected to be in the same order.
        """
        updated_names_weights_dict = dict()

        for key in names_grads_wrt_params_dict.keys():
            # beta = (1 - generated_beta * meta-learned per-step-per-layer bias term)
            # alpha = generated_alpha * meta-learned per-step-per-layer bias term)
            if self.alfa:
                if self.random_init:
                    updated_names_weights_dict[key] = (1 - self.names_beta_dict_per_param[key.replace(".", "-")] * generated_beta_params[key] * self.names_beta_dict[key.replace(".", "-")][num_step]) * names_weights_dict[key] - generated_alpha_params[key] * self.names_alpha_dict[key.replace(".", "-")][num_step] * names_grads_wrt_params_dict[key]
                else:
                    updated_names_weights_dict[key] = (1 - generated_beta_params[key] * self.names_beta_dict[key.replace(".", "-")][num_step]) * names_weights_dict[key] - generated_alpha_params[key] * self.names_alpha_dict[key.replace(".", "-")][num_step] * names_grads_wrt_params_dict[key]
            else:
                updated_names_weights_dict[key] = names_weights_dict[key] - self.init_lr_val * names_grads_wrt_params_dict[key]

        return updated_names_weights_dict

我使用的骨干网络是efficientnet_b4，元学习的网络如下图所示
import torch
from torch import nn
from torch.nn import functional as F
from torch.nn import init
from torch.autograd import Variable
import math
import pdb
from .meta import *
from .utils import (
    round_filters,
    round_repeats,
    drop_connect,
    get_same_padding_conv2d,
    get_same_padding_conv2d_meta,
    get_model_params,
    efficientnet_params,
    load_pretrained_weights,
    Swish,
    MemoryEfficientSwish,
    calculate_output_image_size
)
from network.ScConv import ScConv
class MetaMBConvBlock(MetaModule):


    def __init__(self, block_args, global_params, image_size=None):
        super().__init__()
        self._block_args = block_args
        self._bn_mom = 1 - global_params.batch_norm_momentum # pytorch's difference from tensorflow
        self._bn_eps = global_params.batch_norm_epsilon
        self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1)
        self.id_skip = block_args.id_skip  # whether to use skip connection and drop connect

        # Expansion phase (Inverted Bottleneck)
        inp = self._block_args.input_filters  # number of input channels
        oup = self._block_args.input_filters * self._block_args.expand_ratio  # number of output channels
        if self._block_args.expand_ratio != 1:
            Conv2d = get_same_padding_conv2d_meta(image_size=image_size)
            self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False)
            self._bn0 = MetaBatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)

        # Depthwise convolution phase
        k = self._block_args.kernel_size
        s = self._block_args.stride
        Conv2d = get_same_padding_conv2d_meta(image_size=image_size)
        self._depthwise_conv = Conv2d(
            in_channels=oup, out_channels=oup, groups=oup,  # groups makes it depthwise
            kernel_size=k, stride=s, bias=False)
        self._bn1 = MetaBatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
        
        image_size = calculate_output_image_size(image_size, s)

        # Squeeze and Excitation layer, if desired
        if self.has_se:
            Conv2d = get_same_padding_conv2d_meta(image_size=(1,1))
            num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio))
            self._se_reduce = Conv2d(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1)
            self._se_expand = Conv2d(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1)

        # Pointwise convolution phase
        final_oup = self._block_args.output_filters
        Conv2d = get_same_padding_conv2d_meta(image_size=image_size)
        self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False)
        self._bn2 = MetaBatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps)
        

        self._swish = MemoryEfficientSwish()

    def forward(self, inputs, drop_connect_rate=None):
        """MBConvBlock's forward function.

        Args:
            inputs (tensor): Input tensor.
            drop_connect_rate (bool): Drop connect rate (float, between 0 and 1).

        Returns:
            Output of this block after processing.
        """

        # Expansion and Depthwise Convolution
        x = inputs
        if self._block_args.expand_ratio != 1:
            x = self._expand_conv(inputs)
            x = self._bn0(x)
            x = self._swish(x)

        x = self._depthwise_conv(x)
        x = self._bn1(x)
        
        x = self._swish(x)

        # Squeeze and Excitation
        if self.has_se:
            x_squeezed = F.adaptive_avg_pool2d(x, 1)
            x_squeezed = self._se_reduce(x_squeezed)
            x_squeezed = self._swish(x_squeezed)
            x_squeezed = self._se_expand(x_squeezed)
            x = torch.sigmoid(x_squeezed) * x

        # Pointwise Convolution
        x = self._project_conv(x)
        x = self._bn2(x)
        
        # Skip connection and drop connect
        input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters
        if self.id_skip and self._block_args.stride == 1 and input_filters == output_filters:
            # The combination of skip connection and drop connect brings about stochastic depth.
            if drop_connect_rate:
                x = drop_connect(x, p=drop_connect_rate, training=self.training)
            x = x + inputs  # skip connection
        return x

    def set_swish(self, memory_efficient=True):
        """Sets swish function as memory efficient (for training) or standard (for export).

        Args:
            memory_efficient (bool): Whether to use memory-efficient version of swish.
        """
        self._swish = MemoryEfficientSwish() if memory_efficient else Swish()


class MetaEfficientNet(MetaModule):

    def __init__(self, blocks_args=None, global_params=None,in_channels=3):
        super().__init__()

        assert isinstance(blocks_args, list), 'blocks_args should be a list'
        assert len(blocks_args) > 0, 'block args must be greater than 0'
        self._global_params = global_params
        self._blocks_args = blocks_args

        # Batch norm parameters
        bn_mom = 1 - self._global_params.batch_norm_momentum
        bn_eps = self._global_params.batch_norm_epsilon

        # Get stem static or dynamic convolution depending on image size
        image_size = global_params.image_size
        Conv2d = get_same_padding_conv2d_meta(image_size=image_size)

        # Stem
        in_channels = in_channels  # rgb
        out_channels = round_filters(32, self._global_params)  # number of output channels
        self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False)
        self._bn0 = MetaBatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
        image_size = calculate_output_image_size(image_size, 2)

        # Build blocks
        num_block = len(self._blocks_args)
        self._blocks = nn.ModuleList([])
        
        for i,block_args in enumerate(self._blocks_args):

            # Update block input and output filters based on depth multiplier.
            block_args = block_args._replace(
                input_filters=round_filters(block_args.input_filters, self._global_params),
                output_filters=round_filters(block_args.output_filters, self._global_params),
                num_repeat=round_repeats(block_args.num_repeat, self._global_params)
            )

            # The first block needs to take care of stride and filter size increase.

            self._blocks.append(MetaMBConvBlock(block_args, self._global_params, image_size=image_size))
            #self._blocks.append(MBConvBlockIN(block_args, self._global_params, image_size=image_size))
            image_size = calculate_output_image_size(image_size, block_args.stride)
            if block_args.num_repeat > 1: # modify block_args to keep same output size
                block_args = block_args._replace(input_filters=block_args.output_filters, stride=1)
            for _ in range(block_args.num_repeat - 1):
                self._blocks.append(MetaMBConvBlock(block_args, self._global_params, image_size=image_size))
        # Head
        in_channels = block_args.output_filters  # output of final block
        out_channels = round_filters(1280, self._global_params)
        Conv2d = get_same_padding_conv2d_meta(image_size=image_size)
        self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self._bn1 = MetaBatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)

        # Final linear layer
        self._avg_pooling = nn.AdaptiveAvgPool2d(1)
        self._dropout = nn.Dropout(self._global_params.dropout_rate)
        self._fc = MetaLinear(out_channels, self._global_params.num_classes)
        self._swish = MemoryEfficientSwish()

    def set_swish(self, memory_efficient=True):
        """Sets swish function as memory efficient (for training) or standard (for export).

        Args:
            memory_efficient (bool): Whether to use memory-efficient version of swish.

        """
        self._swish = MemoryEfficientSwish() if memory_efficient else Swish()
        for block in self._blocks:
            block.set_swish(memory_efficient)


    def extract_features(self, inputs):
        """use convolution layer to extract feature .

        Args:
            inputs (tensor): Input tensor.

        Returns:
            Output of the final convolution 
            layer in the efficientnet model.
        """
        # Stem
        x = self._swish(self._bn0(self._conv_stem(inputs)))
        # Blocks
        for idx, block in enumerate(self._blocks):
            drop_connect_rate = self._global_params.drop_connect_rate
            if drop_connect_rate:
                drop_connect_rate *= float(idx) / len(self._blocks) # scale drop connect_rate
            x = block(x, drop_connect_rate=drop_connect_rate)
        
        # Head
        x = self._swish(self._bn1(self._conv_head(x)))

        return x
    

    def forward(self, inputs):
        """EfficientNet's forward function.
           Calls extract_features to extract features, applies final linear layer, and returns logits.

        Args:
            inputs (tensor): Input tensor.

        Returns:
            Output of this model after processing.
        """
        bs = inputs.size(0)

        # Convolution layers
        feature = self.extract_features(inputs)
        # Pooling and final linear layer
        x = self._avg_pooling(feature)
        x = x.view(bs, -1)
        x = self._dropout(x)
        x = self._fc(x)


        return x,feature

    @classmethod
    def from_name(cls, model_name, in_channels=3, **override_params):
        """create an efficientnet model according to name.

        Args:
            model_name (str): Name for efficientnet.
            in_channels (int): Input data's channel number.
            override_params (other key word params): 
                Params to override model's global_params.
                Optional key:
                    'width_coefficient', 'depth_coefficient',
                    'image_size', 'dropout_rate',
                    'num_classes', 'batch_norm_momentum',
                    'batch_norm_epsilon', 'drop_connect_rate',
                    'depth_divisor', 'min_depth'

        Returns:
            An efficientnet model.
        """
        cls._check_model_name_is_valid(model_name)
        blocks_args, global_params = get_model_params(model_name, override_params)
        model = cls(blocks_args, global_params)
        model._change_in_channels(in_channels)
        return model

    @classmethod
    def from_pretrained(cls, model_name, weights_path=None, advprop=False, 
                        in_channels=3, num_classes=1000, **override_params):
        """create an efficientnet model according to name.

        Args:
            model_name (str): Name for efficientnet.
            weights_path (None or str): 
                str: path to pretrained weights file on the local disk.
                None: use pretrained weights downloaded from the Internet.
            advprop (bool): 
                Whether to load pretrained weights
                trained with advprop (valid when weights_path is None).
            in_channels (int): Input data's channel number.
            num_classes (int): 
                Number of categories for classification.
                It controls the output size for final linear layer.
            override_params (other key word params): 
                Params to override model's global_params.
                Optional key:
                    'width_coefficient', 'depth_coefficient',
                    'image_size', 'dropout_rate',
                    'num_classes', 'batch_norm_momentum',
                    'batch_norm_epsilon', 'drop_connect_rate',
                    'depth_divisor', 'min_depth'

        Returns:
            A pretrained efficientnet model.
        """
        model = cls.from_name(model_name, num_classes = num_classes, **override_params)
        load_pretrained_weights(model, model_name, weights_path=weights_path, load_fc=(num_classes == 1000), advprop=advprop)
        model._change_in_channels(in_channels)
        return model

    @classmethod
    def get_image_size(cls, model_name):
        """Get the input image size for a given efficientnet model.

        Args:
            model_name (str): Name for efficientnet.

        Returns:
            Input image size (resolution).
        """
        cls._check_model_name_is_valid(model_name)
        _, _, res, _ = efficientnet_params(model_name)
        return res

    @classmethod
    def _check_model_name_is_valid(cls, model_name):
        """Validates model name. 

        Args:
            model_name (str): Name for efficientnet.

        Returns:
            bool: Is a valid name or not.
        """
        valid_models = ['efficientnet-b'+str(i) for i in range(9)]

        # Support the construction of 'efficientnet-l2' without pretrained weights
        valid_models += ['efficientnet-l2']

        if model_name not in valid_models:
            raise ValueError('model_name should be one of: ' + ', '.join(valid_models))
    
    def _change_in_channels(self, in_channels):
        """Adjust model's first convolution layer to in_channels, if in_channels not equals 3.

        Args:
            in_channels (int): Input data's channel number.
        """
        if in_channels != 3:
            Conv2d = get_same_padding_conv2d(image_size = self._global_params.image_size)
            out_channels = round_filters(32, self._global_params)
            self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False)