Add

class_indices.json

@@ -0,0 +1,7 @@
+{
+    "0": "daisy",
+    "1": "dandelion",
+    "2": "roses",
+    "3": "sunflowers",
+    "4": "tulips"
+}

model.py

@@ -0,0 +1,377 @@
+from collections import OrderedDict
+from functools import partial
+from typing import Callable, Optional
+
+import torch.nn as nn
+import torch
+from torch import Tensor
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    """
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    "Deep Networks with Stochastic Depth", https://arxiv.org/pdf/1603.09382.pdf
+
+    This function is taken from the rwightman.
+    It can be seen here:
+    https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py#L140
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """
+    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    "Deep Networks with Stochastic Depth", https://arxiv.org/pdf/1603.09382.pdf
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class ConvBNAct(nn.Module):
+    def __init__(self,
+                 in_planes: int,
+                 out_planes: int,
+                 kernel_size: int = 3,
+                 stride: int = 1,
+                 groups: int = 1,
+                 norm_layer: Optional[Callable[..., nn.Module]] = None,
+                 activation_layer: Optional[Callable[..., nn.Module]] = None):
+        super(ConvBNAct, self).__init__()
+
+        padding = (kernel_size - 1) // 2
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if activation_layer is None:
+            activation_layer = nn.SiLU  # alias Swish  (torch>=1.7)
+
+        self.conv = nn.Conv2d(in_channels=in_planes,
+                              out_channels=out_planes,
+                              kernel_size=kernel_size,
+                              stride=stride,
+                              padding=padding,
+                              groups=groups,
+                              bias=False)
+
+        self.bn = norm_layer(out_planes)
+        self.act = activation_layer()
+
+    def forward(self, x):
+        result = self.conv(x)
+        result = self.bn(result)
+        result = self.act(result)
+
+        return result
+
+
+class SqueezeExcite(nn.Module):
+    def __init__(self,
+                 input_c: int,   # block input channel
+                 expand_c: int,  # block expand channel
+                 se_ratio: float = 0.25):
+        super(SqueezeExcite, self).__init__()
+        squeeze_c = int(input_c * se_ratio)
+        self.conv_reduce = nn.Conv2d(expand_c, squeeze_c, 1)
+        self.act1 = nn.SiLU()  # alias Swish
+        self.conv_expand = nn.Conv2d(squeeze_c, expand_c, 1)
+        self.act2 = nn.Sigmoid()
+
+    def forward(self, x: Tensor) -> Tensor:
+        scale = x.mean((2, 3), keepdim=True)
+        scale = self.conv_reduce(scale)
+        scale = self.act1(scale)
+        scale = self.conv_expand(scale)
+        scale = self.act2(scale)
+        return scale * x
+
+
+class MBConv(nn.Module):
+    def __init__(self,
+                 kernel_size: int,
+                 input_c: int,
+                 out_c: int,
+                 expand_ratio: int,
+                 stride: int,
+                 se_ratio: float,
+                 drop_rate: float,
+                 norm_layer: Callable[..., nn.Module]):
+        super(MBConv, self).__init__()
+
+        if stride not in [1, 2]:
+            raise ValueError("illegal stride value.")
+
+        self.has_shortcut = (stride == 1 and input_c == out_c)
+
+        activation_layer = nn.SiLU  # alias Swish
+        expanded_c = input_c * expand_ratio
+
+        # 在EfficientNetV2中，MBConv中不存在expansion=1的情况所以conv_pw肯定存在
+        assert expand_ratio != 1
+        # Point-wise expansion
+        self.expand_conv = ConvBNAct(input_c,
+                                     expanded_c,
+                                     kernel_size=1,
+                                     norm_layer=norm_layer,
+                                     activation_layer=activation_layer)
+
+        # Depth-wise convolution
+        self.dwconv = ConvBNAct(expanded_c,
+                                expanded_c,
+                                kernel_size=kernel_size,
+                                stride=stride,
+                                groups=expanded_c,
+                                norm_layer=norm_layer,
+                                activation_layer=activation_layer)
+
+        self.se = SqueezeExcite(input_c, expanded_c, se_ratio) if se_ratio > 0 else nn.Identity()
+
+        # Point-wise linear projection
+        self.project_conv = ConvBNAct(expanded_c,
+                                      out_planes=out_c,
+                                      kernel_size=1,
+                                      norm_layer=norm_layer,
+                                      activation_layer=nn.Identity)  # 注意这里没有激活函数，所有传入Identity
+
+        self.out_channels = out_c
+
+        # 只有在使用shortcut连接时才使用dropout层
+        self.drop_rate = drop_rate
+        if self.has_shortcut and drop_rate > 0:
+            self.dropout = DropPath(drop_rate)
+
+    def forward(self, x: Tensor) -> Tensor:
+        result = self.expand_conv(x)
+        result = self.dwconv(result)
+        result = self.se(result)
+        result = self.project_conv(result)
+
+        if self.has_shortcut:
+            if self.drop_rate > 0:
+                result = self.dropout(result)
+            result += x
+
+        return result
+
+
+class FusedMBConv(nn.Module):
+    def __init__(self,
+                 kernel_size: int,
+                 input_c: int,
+                 out_c: int,
+                 expand_ratio: int,
+                 stride: int,
+                 se_ratio: float,
+                 drop_rate: float,
+                 norm_layer: Callable[..., nn.Module]):
+        super(FusedMBConv, self).__init__()
+
+        assert stride in [1, 2]
+        assert se_ratio == 0
+
+        self.has_shortcut = stride == 1 and input_c == out_c
+        self.drop_rate = drop_rate
+
+        self.has_expansion = expand_ratio != 1
+
+        activation_layer = nn.SiLU  # alias Swish
+        expanded_c = input_c * expand_ratio
+
+        # 只有当expand ratio不等于1时才有expand conv
+        if self.has_expansion:
+            # Expansion convolution
+            self.expand_conv = ConvBNAct(input_c,
+                                         expanded_c,
+                                         kernel_size=kernel_size,
+                                         stride=stride,
+                                         norm_layer=norm_layer,
+                                         activation_layer=activation_layer)
+
+            self.project_conv = ConvBNAct(expanded_c,
+                                          out_c,
+                                          kernel_size=1,
+                                          norm_layer=norm_layer,
+                                          activation_layer=nn.Identity)  # 注意没有激活函数
+        else:
+            # 当只有project_conv时的情况
+            self.project_conv = ConvBNAct(input_c,
+                                          out_c,
+                                          kernel_size=kernel_size,
+                                          stride=stride,
+                                          norm_layer=norm_layer,
+                                          activation_layer=activation_layer)  # 注意有激活函数
+
+        self.out_channels = out_c
+
+        # 只有在使用shortcut连接时才使用dropout层
+        self.drop_rate = drop_rate
+        if self.has_shortcut and drop_rate > 0:
+            self.dropout = DropPath(drop_rate)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.has_expansion:
+            result = self.expand_conv(x)
+            result = self.project_conv(result)
+        else:
+            result = self.project_conv(x)
+
+        if self.has_shortcut:
+            if self.drop_rate > 0:
+                result = self.dropout(result)
+
+            result += x
+
+        return result
+
+
+class EfficientNetV2(nn.Module):
+    def __init__(self,
+                 model_cnf: list,
+                 num_classes: int = 1000,
+                 num_features: int = 1280,
+                 dropout_rate: float = 0.2,
+                 drop_connect_rate: float = 0.2):
+        super(EfficientNetV2, self).__init__()
+
+        for cnf in model_cnf:
+            assert len(cnf) == 8
+
+        norm_layer = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.1)
+
+        stem_filter_num = model_cnf[0][4]
+
+        self.stem = ConvBNAct(3,
+                              stem_filter_num,
+                              kernel_size=3,
+                              stride=2,
+                              norm_layer=norm_layer)  # 激活函数默认是SiLU
+
+        total_blocks = sum([i[0] for i in model_cnf])
+        block_id = 0
+        blocks = []
+        for cnf in model_cnf:
+            repeats = cnf[0]
+            op = FusedMBConv if cnf[-2] == 0 else MBConv
+            for i in range(repeats):
+                blocks.append(op(kernel_size=cnf[1],
+                                 input_c=cnf[4] if i == 0 else cnf[5],
+                                 out_c=cnf[5],
+                                 expand_ratio=cnf[3],
+                                 stride=cnf[2] if i == 0 else 1,
+                                 se_ratio=cnf[-1],
+                                 drop_rate=drop_connect_rate * block_id / total_blocks,
+                                 norm_layer=norm_layer))
+                block_id += 1
+        self.blocks = nn.Sequential(*blocks)
+
+        head_input_c = model_cnf[-1][-3]
+        head = OrderedDict()
+
+        head.update({"project_conv": ConvBNAct(head_input_c,
+                                               num_features,
+                                               kernel_size=1,
+                                               norm_layer=norm_layer)})  # 激活函数默认是SiLU
+
+        head.update({"avgpool": nn.AdaptiveAvgPool2d(1)})
+        head.update({"flatten": nn.Flatten()})
+
+        if dropout_rate > 0:
+            head.update({"dropout": nn.Dropout(p=dropout_rate, inplace=True)})
+        head.update({"classifier": nn.Linear(num_features, num_classes)})
+
+        self.head = nn.Sequential(head)
+
+        # initial weights
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out")
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.ones_(m.weight)
+                nn.init.zeros_(m.bias)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.stem(x)
+        x = self.blocks(x)
+        x = self.head(x)
+
+        return x
+
+
+def efficientnetv2_s(num_classes: int = 1000):
+    """
+    EfficientNetV2
+    https://arxiv.org/abs/2104.00298
+    """
+    # train_size: 300, eval_size: 384
+
+    # repeat, kernel, stride, expansion, in_c, out_c, operator, se_ratio
+    model_config = [[2, 3, 1, 1, 24, 24, 0, 0],
+                    [4, 3, 2, 4, 24, 48, 0, 0],
+                    [4, 3, 2, 4, 48, 64, 0, 0],
+                    [6, 3, 2, 4, 64, 128, 1, 0.25],
+                    [9, 3, 1, 6, 128, 160, 1, 0.25],
+                    [15, 3, 2, 6, 160, 256, 1, 0.25]]
+
+    model = EfficientNetV2(model_cnf=model_config,
+                           num_classes=num_classes,
+                           dropout_rate=0.2)
+    return model
+
+
+def efficientnetv2_m(num_classes: int = 1000):
+    """
+    EfficientNetV2
+    https://arxiv.org/abs/2104.00298
+    """
+    # train_size: 384, eval_size: 480
+
+    # repeat, kernel, stride, expansion, in_c, out_c, operator, se_ratio
+    model_config = [[3, 3, 1, 1, 24, 24, 0, 0],
+                    [5, 3, 2, 4, 24, 48, 0, 0],
+                    [5, 3, 2, 4, 48, 80, 0, 0],
+                    [7, 3, 2, 4, 80, 160, 1, 0.25],
+                    [14, 3, 1, 6, 160, 176, 1, 0.25],
+                    [18, 3, 2, 6, 176, 304, 1, 0.25],
+                    [5, 3, 1, 6, 304, 512, 1, 0.25]]
+
+    model = EfficientNetV2(model_cnf=model_config,
+                           num_classes=num_classes,
+                           dropout_rate=0.3)
+    return model
+
+
+def efficientnetv2_l(num_classes: int = 1000):
+    """
+    EfficientNetV2
+    https://arxiv.org/abs/2104.00298
+    """
+    # train_size: 384, eval_size: 480
+
+    # repeat, kernel, stride, expansion, in_c, out_c, operator, se_ratio
+    model_config = [[4, 3, 1, 1, 32, 32, 0, 0],
+                    [7, 3, 2, 4, 32, 64, 0, 0],
+                    [7, 3, 2, 4, 64, 96, 0, 0],
+                    [10, 3, 2, 4, 96, 192, 1, 0.25],
+                    [19, 3, 1, 6, 192, 224, 1, 0.25],
+                    [25, 3, 2, 6, 224, 384, 1, 0.25],
+                    [7, 3, 1, 6, 384, 640, 1, 0.25]]
+
+    model = EfficientNetV2(model_cnf=model_config,
+                           num_classes=num_classes,
+                           dropout_rate=0.4)
+    return model

my_dataset.py

@@ -0,0 +1,37 @@
+from PIL import Image
+import torch
+from torch.utils.data import Dataset
+
+
+class MyDataSet(Dataset):
+    """自定义数据集"""
+
+    def __init__(self, images_path: list, images_class: list, transform=None):
+        self.images_path = images_path
+        self.images_class = images_class
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.images_path)
+
+    def __getitem__(self, item):
+        img = Image.open(self.images_path[item])
+        # RGB为彩色图片，L为灰度图片
+        if img.mode != 'RGB':
+            raise ValueError("image: {} isn't RGB mode.".format(self.images_path[item]))
+        label = self.images_class[item]
+
+        if self.transform is not None:
+            img = self.transform(img)
+
+        return img, label
+
+    @staticmethod
+    def collate_fn(batch):
+        # 官方实现的default_collate可以参考
+        # https://github.com/pytorch/pytorch/blob/67b7e751e6b5931a9f45274653f4f653a4e6cdf6/torch/utils/data/_utils/collate.py
+        images, labels = tuple(zip(*batch))
+
+        images = torch.stack(images, dim=0)
+        labels = torch.as_tensor(labels)
+        return images, labels

predict.py

@@ -0,0 +1,65 @@
+import os
+import json
+
+import torch
+from PIL import Image
+from torchvision import transforms
+import matplotlib.pyplot as plt
+
+from model import efficientnetv2_m as create_model
+
+
+def main():
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    img_size = {"s": [300, 384],  # train_size, val_size
+                "m": [384, 480],
+                "l": [384, 480]}
+    num_model = "s"
+
+    data_transform = transforms.Compose(
+        [transforms.Resize(img_size[num_model][1]),
+         transforms.CenterCrop(img_size[num_model][1]),
+         transforms.ToTensor(),
+         transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
+
+    # load image
+    img_path = "../d.jpg"
+    assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
+    img = Image.open(img_path)
+    plt.imshow(img)
+    # [N, C, H, W]
+    img = data_transform(img)
+    # expand batch dimension
+    img = torch.unsqueeze(img, dim=0)
+
+    # read class_indict
+    json_path = './class_indices.json'
+    assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
+
+    json_file = open(json_path, "r")
+    class_indict = json.load(json_file)
+
+    # create model
+    model = create_model(num_classes=5).to(device)
+    # load model weights
+    model_weight_path = "./weights/model-20.pth"
+    model.load_state_dict(torch.load(model_weight_path, map_location=device))
+    model.eval()
+    with torch.no_grad():
+        # predict class
+        output = torch.squeeze(model(img.to(device))).cpu()
+        predict = torch.softmax(output, dim=0)
+        predict_cla = torch.argmax(predict).numpy()
+
+    print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
+                                                 predict[predict_cla].numpy())
+    plt.title(print_res)
+    for i in range(len(predict)):
+        print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
+                                                  predict[i].numpy()))
+    plt.show()
+
+
+if __name__ == '__main__':
+    main()

train.py

@@ -0,0 +1,143 @@
+import os
+import math
+import argparse
+
+import torch
+import torch.optim as optim
+from torch.utils.tensorboard import SummaryWriter
+from torchvision import transforms
+import torch.optim.lr_scheduler as lr_scheduler
+
+from model import efficientnetv2_m as create_model
+from my_dataset import MyDataSet
+from utils import read_split_data, train_one_epoch, evaluate
+
+
+def main(args):
+    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+
+    print(args)
+    print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/')
+    tb_writer = SummaryWriter()
+    if os.path.exists("./weights") is False:
+        os.makedirs("./weights")
+
+    train_images_path, train_images_label, val_images_path, val_images_label = read_split_data(args.data_path)
+
+    img_size = {"s": [300, 384],  # train_size, val_size
+                "m": [384, 480],
+                "l": [384, 480]}
+    num_model = "s"
+
+    data_transform = {
+        "train": transforms.Compose([transforms.RandomResizedCrop(img_size[num_model][0]),
+                                     transforms.RandomHorizontalFlip(),
+                                     transforms.ToTensor(),
+                                     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])]),
+        "val": transforms.Compose([transforms.Resize(img_size[num_model][1]),
+                                   transforms.CenterCrop(img_size[num_model][1]),
+                                   transforms.ToTensor(),
+                                   transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])}
+
+    # 实例化训练数据集
+    train_dataset = MyDataSet(images_path=train_images_path,
+                              images_class=train_images_label,
+                              transform=data_transform["train"])
+
+    # 实例化验证数据集
+    val_dataset = MyDataSet(images_path=val_images_path,
+                            images_class=val_images_label,
+                            transform=data_transform["val"])
+
+    batch_size = args.batch_size
+    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
+    print('Using {} dataloader workers every process'.format(nw))
+    train_loader = torch.utils.data.DataLoader(train_dataset,
+                                               batch_size=batch_size,
+                                               shuffle=True,
+                                               pin_memory=True,
+                                               num_workers=nw,
+                                               collate_fn=train_dataset.collate_fn)
+
+    val_loader = torch.utils.data.DataLoader(val_dataset,
+                                             batch_size=batch_size,
+                                             shuffle=False,
+                                             pin_memory=True,
+                                             num_workers=nw,
+                                             collate_fn=val_dataset.collate_fn)
+
+    # 如果存在预训练权重则载入
+    model = create_model(num_classes=args.num_classes).to(device)
+    if args.weights != "":
+        if os.path.exists(args.weights):
+            weights_dict = torch.load(args.weights, map_location=device)
+            load_weights_dict = {k: v for k, v in weights_dict.items()
+                                 if model.state_dict()[k].numel() == v.numel()}
+            print(model.load_state_dict(load_weights_dict, strict=False))
+        else:
+            raise FileNotFoundError("not found weights file: {}".format(args.weights))
+
+    # 是否冻结权重
+    if args.freeze_layers:
+        for name, para in model.named_parameters():
+            # 除head外，其他权重全部冻结
+            if "head" not in name:
+                para.requires_grad_(False)
+            else:
+                print("training {}".format(name))
+
+    pg = [p for p in model.parameters() if p.requires_grad]
+    optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=1E-4)
+    # Scheduler https://arxiv.org/pdf/1812.01187.pdf
+    lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf  # cosine
+    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
+
+    for epoch in range(args.epochs):
+        # train
+        train_loss, train_acc = train_one_epoch(model=model,
+                                                optimizer=optimizer,
+                                                data_loader=train_loader,
+                                                device=device,
+                                                epoch=epoch)
+
+        scheduler.step()
+
+        # validate
+        val_loss, val_acc = evaluate(model=model,
+                                     data_loader=val_loader,
+                                     device=device,
+                                     epoch=epoch)
+
+        tags = ["train_loss", "train_acc", "val_loss", "val_acc", "learning_rate"]
+        tb_writer.add_scalar(tags[0], train_loss, epoch)
+        tb_writer.add_scalar(tags[1], train_acc, epoch)
+        tb_writer.add_scalar(tags[2], val_loss, epoch)
+        tb_writer.add_scalar(tags[3], val_acc, epoch)
+        tb_writer.add_scalar(tags[4], optimizer.param_groups[0]["lr"], epoch)
+
+        torch.save(model.state_dict(), "./weights/model-{}.pth".format(epoch))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num_classes', type=int, default=5)
+    parser.add_argument('--epochs', type=int, default=30)
+    parser.add_argument('--batch-size', type=int, default=8)
+    parser.add_argument('--lr', type=float, default=0.01)
+    parser.add_argument('--lrf', type=float, default=0.01)
+
+    # 数据集所在根目录
+    # http://download.tensorflow.org/example_images/flower_photos.tgz
+    parser.add_argument('--data-path', type=str,
+                        default="../../data_set/flower_data/flower_photos")
+
+    # download model weights
+    # 链接: https://pan.baidu.com/s/1uZX36rvrfEss-JGj4yfzbQ  密码: 5gu1
+    parser.add_argument('--weights', type=str, default='./pre_efficientnetv2-m.pth',
+                        help='initial weights path')
+    parser.add_argument('--freeze-layers', type=bool, default=True)
+    parser.add_argument('--device', default='cuda:0', help='device id (i.e. 0 or 0,1 or cpu)')
+
+    opt = parser.parse_args()
+
+    main(opt)

trans_effv2_weights.py

@@ -0,0 +1,160 @@
+import tensorflow as tf
+import torch
+import numpy as np
+
+
+def main(model_name: str = "efficientnetv2-s",
+         tf_weights_path: str = "./efficientnetv2-s/model",
+         stage0_num: int = 2,
+         fused_conv_num: int = 10):
+
+    except_var = ["global_step"]
+
+    new_weights = {}
+    var_list = [i for i in tf.train.list_variables(tf_weights_path) if "Exponential" not in i[0]]
+    reader = tf.train.load_checkpoint(tf_weights_path)
+    for v in var_list:
+        if v[0] in except_var:
+            continue
+        new_name = v[0].replace(model_name + "/", "").replace("/", ".")
+
+        if "stem" in v[0]:
+            new_name = new_name.replace("conv2d.kernel",
+                                        "conv.weight")
+
+            new_name = new_name.replace("tpu_batch_normalization.beta",
+                                        "bn.bias")
+            new_name = new_name.replace("tpu_batch_normalization.gamma",
+                                        "bn.weight")
+            new_name = new_name.replace("tpu_batch_normalization.moving_mean",
+                                        "bn.running_mean")
+            new_name = new_name.replace("tpu_batch_normalization.moving_variance",
+                                        "bn.running_var")
+        elif "head" in v[0]:
+            new_name = new_name.replace("conv2d.kernel",
+                                        "project_conv.conv.weight")
+            new_name = new_name.replace("dense.kernel",
+                                        "classifier.weight")
+            new_name = new_name.replace("dense.bias",
+                                        "classifier.bias")
+
+            new_name = new_name.replace("tpu_batch_normalization.beta",
+                                        "project_conv.bn.bias")
+            new_name = new_name.replace("tpu_batch_normalization.gamma",
+                                        "project_conv.bn.weight")
+            new_name = new_name.replace("tpu_batch_normalization.moving_mean",
+                                        "project_conv.bn.running_mean")
+            new_name = new_name.replace("tpu_batch_normalization.moving_variance",
+                                        "project_conv.bn.running_var")
+        elif "blocks" in v[0]:
+            # e.g. blocks_0.conv2d.kernel -> 0
+            blocks_id = new_name.split(".", maxsplit=1)[0].replace("blocks_", "")
+            new_name = new_name.replace("blocks_{}".format(blocks_id),
+                                        "blocks.{}".format(blocks_id))
+
+            if int(blocks_id) <= stage0_num - 1:  # expansion=1 fused_mbconv
+                new_name = new_name.replace("conv2d.kernel",
+                                            "project_conv.conv.weight")
+                new_name = new_name.replace("tpu_batch_normalization.beta",
+                                            "project_conv.bn.bias")
+                new_name = new_name.replace("tpu_batch_normalization.gamma",
+                                            "project_conv.bn.weight")
+                new_name = new_name.replace("tpu_batch_normalization.moving_mean",
+                                            "project_conv.bn.running_mean")
+                new_name = new_name.replace("tpu_batch_normalization.moving_variance",
+                                            "project_conv.bn.running_var")
+            else:
+                new_name = new_name.replace("blocks.{}.conv2d.kernel".format(blocks_id),
+                                            "blocks.{}.expand_conv.conv.weight".format(blocks_id))
+                new_name = new_name.replace("tpu_batch_normalization.beta",
+                                            "expand_conv.bn.bias")
+                new_name = new_name.replace("tpu_batch_normalization.gamma",
+                                            "expand_conv.bn.weight")
+                new_name = new_name.replace("tpu_batch_normalization.moving_mean",
+                                            "expand_conv.bn.running_mean")
+                new_name = new_name.replace("tpu_batch_normalization.moving_variance",
+                                            "expand_conv.bn.running_var")
+
+                if int(blocks_id) <= fused_conv_num - 1:  # fused_mbconv
+                    new_name = new_name.replace("blocks.{}.conv2d_1.kernel".format(blocks_id),
+                                                "blocks.{}.project_conv.conv.weight".format(blocks_id))
+                    new_name = new_name.replace("tpu_batch_normalization_1.beta",
+                                                "project_conv.bn.bias")
+                    new_name = new_name.replace("tpu_batch_normalization_1.gamma",
+                                                "project_conv.bn.weight")
+                    new_name = new_name.replace("tpu_batch_normalization_1.moving_mean",
+                                                "project_conv.bn.running_mean")
+                    new_name = new_name.replace("tpu_batch_normalization_1.moving_variance",
+                                                "project_conv.bn.running_var")
+                else:  # mbconv
+                    new_name = new_name.replace("blocks.{}.conv2d_1.kernel".format(blocks_id),
+                                                "blocks.{}.project_conv.conv.weight".format(blocks_id))
+
+                    new_name = new_name.replace("depthwise_conv2d.depthwise_kernel",
+                                                "dwconv.conv.weight")
+
+                    new_name = new_name.replace("tpu_batch_normalization_1.beta",
+                                                "dwconv.bn.bias")
+                    new_name = new_name.replace("tpu_batch_normalization_1.gamma",
+                                                "dwconv.bn.weight")
+                    new_name = new_name.replace("tpu_batch_normalization_1.moving_mean",
+                                                "dwconv.bn.running_mean")
+                    new_name = new_name.replace("tpu_batch_normalization_1.moving_variance",
+                                                "dwconv.bn.running_var")
+
+                    new_name = new_name.replace("tpu_batch_normalization_2.beta",
+                                                "project_conv.bn.bias")
+                    new_name = new_name.replace("tpu_batch_normalization_2.gamma",
+                                                "project_conv.bn.weight")
+                    new_name = new_name.replace("tpu_batch_normalization_2.moving_mean",
+                                                "project_conv.bn.running_mean")
+                    new_name = new_name.replace("tpu_batch_normalization_2.moving_variance",
+                                                "project_conv.bn.running_var")
+
+                    new_name = new_name.replace("se.conv2d.bias",
+                                                "se.conv_reduce.bias")
+                    new_name = new_name.replace("se.conv2d.kernel",
+                                                "se.conv_reduce.weight")
+                    new_name = new_name.replace("se.conv2d_1.bias",
+                                                "se.conv_expand.bias")
+                    new_name = new_name.replace("se.conv2d_1.kernel",
+                                                "se.conv_expand.weight")
+        else:
+            print("not recognized name: " + v[0])
+
+        var = reader.get_tensor(v[0])
+        new_var = var
+        if "conv" in new_name and "weight" in new_name and "bn" not in new_name and "dw" not in new_name:
+            assert len(var.shape) == 4
+            # conv kernel [h, w, c, n] -> [n, c, h, w]
+            new_var = np.transpose(var, (3, 2, 0, 1))
+        elif "bn" in new_name:
+            pass
+        elif "dwconv" in new_name and "weight" in new_name:
+            # dw_kernel [h, w, n, c] -> [n, c, h, w]
+            assert len(var.shape) == 4
+            new_var = np.transpose(var, (2, 3, 0, 1))
+        elif "classifier" in new_name and "weight" in new_name:
+            assert len(var.shape) == 2
+            new_var = np.transpose(var, (1, 0))
+
+        new_weights[new_name] = torch.as_tensor(new_var)
+
+    torch.save(new_weights, "pre_" + model_name + ".pth")
+
+
+if __name__ == '__main__':
+    main(model_name="efficientnetv2-s",
+         tf_weights_path="./efficientnetv2-s/model",
+         stage0_num=2,
+         fused_conv_num=10)
+
+    # main(model_name="efficientnetv2-m",
+    #      tf_weights_path="./efficientnetv2-m/model",
+    #      stage0_num=3,
+    #      fused_conv_num=13)
+
+    # main(model_name="efficientnetv2-l",
+    #      tf_weights_path="./efficientnetv2-l/model",
+    #      stage0_num=4,
+    #      fused_conv_num=18)

utils.py

@@ -0,0 +1,175 @@
+import os
+import sys
+import json
+import pickle
+import random
+
+import torch
+from tqdm import tqdm
+
+import matplotlib.pyplot as plt
+
+
+def read_split_data(root: str, val_rate: float = 0.2):
+    random.seed(0)  # 保证随机结果可复现
+    assert os.path.exists(root), "dataset root: {} does not exist.".format(root)
+
+    # 遍历文件夹，一个文件夹对应一个类别
+    flower_class = [cla for cla in os.listdir(root) if os.path.isdir(os.path.join(root, cla))]
+    # 排序，保证顺序一致
+    flower_class.sort()
+    # 生成类别名称以及对应的数字索引
+    class_indices = dict((k, v) for v, k in enumerate(flower_class))
+    json_str = json.dumps(dict((val, key) for key, val in class_indices.items()), indent=4)
+    with open('class_indices.json', 'w') as json_file:
+        json_file.write(json_str)
+
+    train_images_path = []  # 存储训练集的所有图片路径
+    train_images_label = []  # 存储训练集图片对应索引信息
+    val_images_path = []  # 存储验证集的所有图片路径
+    val_images_label = []  # 存储验证集图片对应索引信息
+    every_class_num = []  # 存储每个类别的样本总数
+    supported = [".jpg", ".JPG", ".png", ".PNG"]  # 支持的文件后缀类型
+    # 遍历每个文件夹下的文件
+    for cla in flower_class:
+        cla_path = os.path.join(root, cla)
+        # 遍历获取supported支持的所有文件路径
+        images = [os.path.join(root, cla, i) for i in os.listdir(cla_path)
+                  if os.path.splitext(i)[-1] in supported]
+        # 获取该类别对应的索引
+        image_class = class_indices[cla]
+        # 记录该类别的样本数量
+        every_class_num.append(len(images))
+        # 按比例随机采样验证样本
+        val_path = random.sample(images, k=int(len(images) * val_rate))
+
+        for img_path in images:
+            if img_path in val_path:  # 如果该路径在采样的验证集样本中则存入验证集
+                val_images_path.append(img_path)
+                val_images_label.append(image_class)
+            else:  # 否则存入训练集
+                train_images_path.append(img_path)
+                train_images_label.append(image_class)
+
+    print("{} images were found in the dataset.".format(sum(every_class_num)))
+    print("{} images for training.".format(len(train_images_path)))
+    print("{} images for validation.".format(len(val_images_path)))
+
+    plot_image = False
+    if plot_image:
+        # 绘制每种类别个数柱状图
+        plt.bar(range(len(flower_class)), every_class_num, align='center')
+        # 将横坐标0,1,2,3,4替换为相应的类别名称
+        plt.xticks(range(len(flower_class)), flower_class)
+        # 在柱状图上添加数值标签
+        for i, v in enumerate(every_class_num):
+            plt.text(x=i, y=v + 5, s=str(v), ha='center')
+        # 设置x坐标
+        plt.xlabel('image class')
+        # 设置y坐标
+        plt.ylabel('number of images')
+        # 设置柱状图的标题
+        plt.title('flower class distribution')
+        plt.show()
+
+    return train_images_path, train_images_label, val_images_path, val_images_label
+
+
+def plot_data_loader_image(data_loader):
+    batch_size = data_loader.batch_size
+    plot_num = min(batch_size, 4)
+
+    json_path = './class_indices.json'
+    assert os.path.exists(json_path), json_path + " does not exist."
+    json_file = open(json_path, 'r')
+    class_indices = json.load(json_file)
+
+    for data in data_loader:
+        images, labels = data
+        for i in range(plot_num):
+            # [C, H, W] -> [H, W, C]
+            img = images[i].numpy().transpose(1, 2, 0)
+            # 反Normalize操作
+            img = (img * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]) * 255
+            label = labels[i].item()
+            plt.subplot(1, plot_num, i+1)
+            plt.xlabel(class_indices[str(label)])
+            plt.xticks([])  # 去掉x轴的刻度
+            plt.yticks([])  # 去掉y轴的刻度
+            plt.imshow(img.astype('uint8'))
+        plt.show()
+
+
+def write_pickle(list_info: list, file_name: str):
+    with open(file_name, 'wb') as f:
+        pickle.dump(list_info, f)
+
+
+def read_pickle(file_name: str) -> list:
+    with open(file_name, 'rb') as f:
+        info_list = pickle.load(f)
+        return info_list
+
+
+def train_one_epoch(model, optimizer, data_loader, device, epoch):
+    model.train()
+    loss_function = torch.nn.CrossEntropyLoss()
+    accu_loss = torch.zeros(1).to(device)  # 累计损失
+    accu_num = torch.zeros(1).to(device)   # 累计预测正确的样本数
+    optimizer.zero_grad()
+
+    sample_num = 0
+    data_loader = tqdm(data_loader)
+    for step, data in enumerate(data_loader):
+        images, labels = data
+        sample_num += images.shape[0]
+
+        pred = model(images.to(device))
+        pred_classes = torch.max(pred, dim=1)[1]
+        accu_num += torch.eq(pred_classes, labels.to(device)).sum()
+
+        loss = loss_function(pred, labels.to(device))
+        loss.backward()
+        accu_loss += loss.detach()
+
+        data_loader.desc = "[train epoch {}] loss: {:.3f}, acc: {:.3f}".format(epoch,
+                                                                               accu_loss.item() / (step + 1),
+                                                                               accu_num.item() / sample_num)
+
+        if not torch.isfinite(loss):
+            print('WARNING: non-finite loss, ending training ', loss)
+            sys.exit(1)
+
+        optimizer.step()
+        optimizer.zero_grad()
+
+    return accu_loss.item() / (step + 1), accu_num.item() / sample_num
+
+
+@torch.no_grad()
+def evaluate(model, data_loader, device, epoch):
+    loss_function = torch.nn.CrossEntropyLoss()
+
+    model.eval()
+
+    accu_num = torch.zeros(1).to(device)   # 累计预测正确的样本数
+    accu_loss = torch.zeros(1).to(device)  # 累计损失
+
+    sample_num = 0
+    data_loader = tqdm(data_loader)
+    for step, data in enumerate(data_loader):
+        images, labels = data
+        sample_num += images.shape[0]
+
+        pred = model(images.to(device))
+        pred_classes = torch.max(pred, dim=1)[1]
+        accu_num += torch.eq(pred_classes, labels.to(device)).sum()
+
+        loss = loss_function(pred, labels.to(device))
+        accu_loss += loss
+
+        data_loader.desc = "[valid epoch {}] loss: {:.3f}, acc: {:.3f}".format(epoch,
+                                                                               accu_loss.item() / (step + 1),
+                                                                               accu_num.item() / sample_num)
+
+    return accu_loss.item() / (step + 1), accu_num.item() / sample_num

weights/model-20.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e3027cc78d0448540d99ed074391d48031c1ab3c6d23e3868bb83a7c9c90c9e
+size 213027833