下面的是我训练的代码，我想在其中加入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

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):
    """
    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)
    output = post_function(output)
    output = output.squeeze()
    
    prediction = (output >= 0.5).float()
    if feat==False:
        return prediction, output
    else:
        return prediction, output,feature


def train(model,optimizer,train_dataloader,meta_dataloader,criterion_oc,epoch,epoch_size,device):
    '''
    train the LTW framework.
    '''
    if meta_dataloader!=None:
        meta_dataloader_iter = iter(meta_dataloader)
    losses = AverageMeter()
    acces = AverageMeter()
    meta_losses = AverageMeter()
    meta_acces = AverageMeter()
    post_function=nn.Sigmoid()
    criterion_norm = torch.nn.BCELoss().to(device)
    criterion = torch.nn.BCELoss(reduce = False).to(device)
    meta_lr = metalr * ((gamma ** int(epoch >= step_size)) * (gamma ** int(epoch >= step_size*2)) * (gamma ** int(epoch >= step_size*3)) * (gamma ** int(epoch >= step_size*4)))         
    
    print(f"meta lr = {meta_lr}")  

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

        meta_model = model_selection(model_name=model_name, num_classes=1, pretrained = False)
        #to get theta'
        if parallel:
            meta_model.copyModel(model.module)
        else:
            meta_model.copyModel(model)
        prediction, output,feature = model_forward(images,model,feat = True)
        compact_loss = criterion_oc(output,targets)

        cost = criterion(output, targets)
        cost_v = torch.reshape(cost, (len(cost), 1))

        # f_lambda = fnet(feature)

        # f_lambda_norm = nn.Sigmoid()(f_lambda)

        l_f_model = torch.sum(cost_v)/len(cost_v) + lamda * compact_loss

        
        model.zero_grad()
        if parallel:
            grads = torch.autograd.grad(l_f_model, (model.module.params()),create_graph=True)
        else:
            grads = torch.autograd.grad(l_f_model, (model.params()),create_graph=True)
        #Virtual update model
        meta_model.update_params(lr_inner=meta_lr, source_params=grads,solver = 'adam')
        
        # del grads
        if meta_dataloader!=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)
            # meta_p = meta_datas[3]
        else:
            #use opposite image
            # opposite_images = datas[4].to(device)
            # opposite_targets = (targets-1)*-1
            # opposite_p = list(datas[5])
            # domain_images = datas[6].to(device)
            # domain_targets = targets
            # domain_p = list(datas[7])


            # meta_images = opposite_images
            # meta_targets = opposite_targets
            meta_images = images
            meta_targets = targets
        
        prediction_meta, output_meta = model_forward(meta_images,meta_model)
        l_g_meta = criterion_norm(output_meta, meta_targets)
        # with torch.no_grad():
        #     w_new = fnet(feature)
        #     w_new_norm = nn.Sigmoid()(w_new)

        # loss = torch.sum(cost_v * (w_new_norm+1))/len(cost_v)
        loss = torch.sum(cost_v)/len(cost_v)
        loss_add = loss + alpha*l_g_meta + lamda*compact_loss

        optimizer.zero_grad()
        # optimizer_fnet.zero_grad()
        loss_add.backward()
        optimizer.step()
        # optimizer_fnet.step()     
        
        acc = (prediction==targets).float().mean()
        meta_acc = (prediction_meta==meta_targets).float().mean()

        meta_acces.update(meta_acc.item(), targets.size(0))
        acces.update(acc.item(), targets.size(0))

        losses.update(loss.item(), targets.size(0))
        meta_losses.update(l_g_meta.item(),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}) || Accuracy:{acces.val:.4f}({acces.avg:.4f}) || MetaAccuracy:{meta_acces.val:.4f}({meta_acces.avg:.4f}) || 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)

    
    if parallel:
        optimizer = torch.optim.Adam(model.module.params(), lr=learning_rate, betas=(beta1,beta2))
    else:
        optimizer = torch.optim.Adam(model.params(), lr=learning_rate, betas=(beta1,beta2))

    # optimizer_fnet = torch.optim.Adam(fnet.params(), lr=plr, betas=(beta1,beta2))

    scheduler = lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)
    # scheduler_fnet = lr_scheduler.StepLR(optimizer_fnet, 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)
        #train2(model,optimizer,train_dataloader,criterion,epoch,epoch_size,device,meta_dataloader=None)

        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

我使用的网络模型如下面所示from __future__ import absolute_import

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
from .meta_efficientnet.meta import *
import pdb
import pretrainedmodels
import torchvision

class MetaTranserModel(MetaModule):
    def __init_with_imagenet(self, baseModel):
        from .efficientnet_pytorch import EfficientNet
        model = EfficientNet.from_pretrained(self.model_name)
        baseModel.copyWeight(model.state_dict())

    def getBase(self,pretrained):
        from .meta_efficientnet import MetaEfficientNet
        baseModel = MetaEfficientNet.from_name(self.model_name)
        if pretrained:
            self.__init_with_imagenet(baseModel)
        return baseModel

    def __init__(self, model_name, num_classes=1,pretrained = True):
        super(MetaTranserModel, self).__init__()
        self.model_name = model_name
        if model_name.startswith('efficientnet'):
            self.base = self.getBase(pretrained)
            self.num_ftrs = self.base._fc.in_features
            self.base._fc = MetaLinear(self.num_ftrs, num_classes)

    def forward(self, x, params=None):
        output,feature = self.base(x, params=params)
        return output,feature

class FNet(MetaModule):
    def __init__(self,in_channel):
        super(FNet, self).__init__()
        def conv_dw(inp, oup, stride):
            return nn.Sequential(
                MetaConv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
                MetaBatchNorm2d(inp),
                nn.ReLU(inplace=True),
                MetaConv2d(inp, oup, 1, 1, 0, bias=False),
                MetaBatchNorm2d(oup),
                nn.ReLU(inplace=True),
            )
        self.feature = nn.Sequential(
            conv_dw(in_channel,in_channel//2,1),
            conv_dw(in_channel//2,in_channel//4,1),
        )
        self._avg_pooling = nn.AdaptiveAvgPool2d(1)
        self._dropout = nn.Dropout(0.5)
        self.bn = MetaBatchNorm1d(1)
        self.fc = MetaLinear(in_channel//4, 1)

    def forward(self, x):
        feature = self.feature(x)
        x = self._avg_pooling(feature)
        x = x.view(x.size(0), -1)
        x = self._dropout(x)
        output = self.fc(x)
        output = self.bn(output)
        return output

if __name__ == '__main__':
    model = MetaTranserModel("efficientnet-b0",0)

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 collections import OrderedDict

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, params=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.
        """
        param_dict = params if params is not None else dict()
        # Expansion and Depthwise Convolution
        x = inputs
        if self._block_args.expand_ratio != 1:
            x = self._expand_conv(x, params=self.get_subdict(param_dict, '_expand_conv'))
            x = self._bn0(x, params=self.get_subdict(param_dict, '_bn0'))
            x = self._swish(x)

        x = self._depthwise_conv(x, params=self.get_subdict(param_dict, '_depthwise_conv'))
        x = self._bn1(x, params=self.get_subdict(param_dict, '_bn1'))
        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, params=self.get_subdict(param_dict, '_se_reduce'))
            x_squeezed = self._swish(x_squeezed)
            x_squeezed = self._se_expand(x_squeezed, params=self.get_subdict(param_dict, '_se_expand'))
            x = torch.sigmoid(x_squeezed) * x

        # Pointwise Convolution
        x = self._project_conv(x, params=self.get_subdict(param_dict, '_project_conv'))
        x = self._bn2(x, params=self.get_subdict(param_dict, '_bn2'))
        
        # 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, params):
        if params is not None:
            print(f"Keys in params: {list(params.keys())}")
        else:
            print("params is None")
        # Stem
        weight = params['conv_stem.weight']
        x = F.conv2d(inputs, weight, bias=None, stride=self._conv_stem.stride, padding=self._conv_stem.padding)
        # Batch Norm 0
        bn_weight = params['bn0.weight']
        bn_bias = params['bn0.bias']
        x = F.batch_norm(
            x,
            self._bn0.running_mean,
            self._bn0.running_var,
            weight=bn_weight,
            bias=bn_bias,
            training=self.training,
            momentum=self._bn0.momentum,
            eps=self._bn0.eps
        )
        x = self._swish(x)

        # Blocks
        for idx, block in enumerate(self._blocks):
            x = block(x, params, prefix=f'_blocks.{idx}')

        # Head
        weight = params['conv_head.weight']
        x = F.conv2d(x, weight, bias=None, stride=self._conv_head.stride, padding=self._conv_head.padding)
        # Batch Norm 1
        bn_weight = params['bn1.weight']
        bn_bias = params['bn1.bias']
        x = F.batch_norm(
            x,
            self._bn1.running_mean,
            self._bn1.running_var,
            weight=bn_weight,
            bias=bn_bias,
            training=self.training,
            momentum=self._bn1.momentum,
            eps=self._bn1.eps
        )
        x = self._swish(x)
        return x
    

    def forward(self, inputs, params=None):
        """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)
        if params is None:
            params = OrderedDict(self.named_parameters())
        # Convolution layers
        feature = self.extract_features(inputs, params=params)
        # Pooling and final linear layer
        x = self._avg_pooling(feature)
        x = x.view(bs, -1)
        x = self._dropout(x)
        weight = params['fc.weight']
        bias = params.get('fc.bias', None)
        x = F.linear(x, weight, bias)

        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)