{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":7327,"databundleVersionId":861871,"sourceType":"competition"}],"dockerImageVersionId":30747,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 微调Swin-B完成图像分类任务 (ImageNet-Dog Breeds)\n# Fine-tuning Swin-B for Image Classification Task (ImageNet-Dog Breeds)","metadata":{}},{"cell_type":"markdown","source":"# What‘s new: \n1. fixed the bugs in training new head and broaden the choices of different Swin Transformer like Sin-T & Swin-B.   \n2.I tested blank Swin-T on a RTX4090, there is still no enough memory for all gradients.  \n3.I also tested more unfrozen layers in finetuning but the precision didn't increase.\n4. Using Kaiming initialization in the new head slightly increased precision.\n5.The score can reach 0.27x, but it's not possible to directly complete and submit the task. You need to download it from the 'output' folder and then submit it. I don't know why.  \n6.I suspect that there might be some issues with the official dataset and scoring system from seven years ago. When I decompressed the files, many images prompted header file errors, and the scores could not be reduced.\n\n# 更新内容:\n1.修复了训练新头的bug，增加了Swin-T，Swin-B的实现。  \n2.我在单张RTX4090上测试了空白的Swin-T，仍然没有足够的内存供所有梯度反向传播使用。  \n3.使用恺明初始化能略微提升精度  \n4.我还在微调中测试了更多未冻结层，但精度没有提高，果然还是要相信前人智慧  \n5.评分可以到0.27x，但是没办法直接跑完提交，需要在output里下载然后提交   \n6.我怀疑7年前的官方数据集和评分系统可能有一点问题，解压的时候很多图片会提示头文件错误，分数也降不下去","metadata":{}},{"cell_type":"markdown","source":"#### 1.部署和可视化参考了李沐老师的d2l库，为方便在线使用重新实现了库中的函数\n#### 2.笔记本中包含了完整的Swin Transformer架构的手动实现，调用时选择pretrain=False即可从头训练\n#### 3.进一步了解 Swin Transformer：https://github.com/microsoft/Swin-Transformer & https://arxiv.org/pdf/2103.14030\n#### 4.警告：经测试，Kaggle在线GPU仅勉强能够微调Swin Transformer，batch_size = 128时显存不足以从头开始训练Swin-S。从零开始构建请仿照swin_small_patch4_window7_224()等函数修改参数","metadata":{}},{"cell_type":"markdown","source":"#### The deployment and visualization refer to the d2l library by M. Li. For easy online use, the functions in the library have been re-implemented.\n#### The notebook contains a complete manual implementation of the Swin Transformer architecture. Select pretrain=False when calling to train from scratch.\n#### To learn more about Swin Transformer: https://github.com/microsoft/Swin-Transformer & https://arxiv.org/pdf/2103.14030\n#### Warning: Tested, the online GPU on Kaggle can barely used for fine-tuning Swin-S. The memory is insufficient to train Swin-S from scratch while batch_size = 128. If you want to train from 0, please use function swin_small_patch4_window7_224() and another 2 structures.","metadata":{}},{"cell_type":"code","source":"import collections\nimport math\nimport os\nimport shutil\nimport re\nimport hashlib\nimport shutil\nimport requests\nimport time\nfrom collections import defaultdict\nimport zipfile\nimport tarfile\n\nimport numpy as np\nfrom PIL import Image\n%matplotlib inline\nfrom matplotlib_inline import backend_inline\nfrom matplotlib import pyplot as plt\nimport pandas as pd\nimport torch\nimport torchvision\nfrom torch import nn\nfrom torch.nn import functional as F\nimport torch.utils.checkpoint as checkpoint\nfrom typing import Optional\nfrom IPython import display\nfrom torch.utils import data\nfrom torchvision import transforms\nfrom torchvision.models.swin_transformer import swin_s, Swin_S_Weights, swin_t, Swin_T_Weights, swin_b, Swin_B_Weights","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:29:07.029409Z","iopub.execute_input":"2024-08-18T08:29:07.029722Z","iopub.status.idle":"2024-08-18T08:29:10.308624Z","shell.execute_reply.started":"2024-08-18T08:29:07.029690Z","shell.execute_reply":"2024-08-18T08:29:10.307840Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 封装基本的绘图和计算工具（from d2l）\n### Preparing some data viz and computing rudimentary methods","metadata":{}},{"cell_type":"code","source":"def use_svg_display(): #@save\n    \"\"\"使用svg格式在Jupyter中显示绘图\"\"\"\n    backend_inline.set_matplotlib_formats('svg')\n\ndef set_axes(axes, xlabel, ylabel, xlim, ylim, xscale, yscale, legend):\n    \"\"\"Set the axes for matplotlib.\"\"\"\n    axes.set_xlabel(xlabel)\n    axes.set_ylabel(ylabel)\n    axes.set_xscale(xscale)\n    axes.set_yscale(yscale)\n    axes.set_xlim(xlim)\n    axes.set_ylim(ylim)\n    if legend:\n        axes.legend(legend)\n    axes.grid()\n\ndef set_figsize(figsize=(3.5, 2.5)): #@save\n    \"\"\"设置matplotlib的图表大小\"\"\"\n    use_svg_display()\n    plt.rcParams['figure.figsize'] = figsize\n\ndef plot(X, Y=None, xlabel=None, ylabel=None, legend=None, xlim=None,\n         ylim=None, xscale='linear', yscale='linear',\n         fmts=('-', 'm--', 'g-.', 'r:'), figsize=(3.5, 2.5), axes=None):\n    \"\"\"Plot data points.\"\"\"\n    if legend is None:\n        legend = []\n\n    set_figsize(figsize)\n    axes = axes if axes else plt.gca()\n\n    # Return True if `X` (tensor or list) has 1 axis\n    def has_one_axis(X):\n        return (hasattr(X, \"ndim\") and X.ndim == 1 or\n                isinstance(X, list) and not hasattr(X[0], \"__len__\"))\n\n    if has_one_axis(X):\n        X = [X]\n    if Y is None:\n        X, Y = [[]] * len(X), X\n    elif has_one_axis(Y):\n        Y = [Y]\n    if len(X) != len(Y):\n        X = X * len(Y)\n    axes.cla()\n    for x, y, fmt in zip(X, Y, fmts):\n        if len(x):\n            axes.plot(x, y, fmt)\n        else:\n            axes.plot(y, fmt)\n    set_axes(axes, xlabel, ylabel, xlim, ylim, xscale, yscale, legend)\n    \nclass Accumulator:\n    \"\"\"For accumulating sums over `n` variables.\"\"\"\n    def __init__(self, n):\n        self.data = [0.0] * n\n\n    def add(self, *args):\n        self.data = [a + float(b) for a, b in zip(self.data, args)]\n\n    def reset(self):\n        self.data = [0.0] * len(self.data)\n\n    def __getitem__(self, idx):\n        return self.data[idx]\n    \nclass Animator:\n    \"\"\"For plotting data in animation.\"\"\"\n    def __init__(self, xlabel=None, ylabel=None, legend=None, xlim=None,\n                 ylim=None, xscale='linear', yscale='linear',\n                 fmts=('-', 'm--', 'g-.', 'r:'), nrows=1, ncols=1,\n                 figsize=(3.5, 2.5)):\n        # Incrementally plot multiple lines\n        if legend is None:\n            legend = []\n        use_svg_display()\n        self.fig, self.axes = plt.subplots(nrows, ncols, figsize=figsize)\n        if nrows * ncols == 1:\n            self.axes = [self.axes,]\n        # Use a lambda function to capture arguments\n        self.config_axes = lambda: set_axes(self.axes[\n            0], xlabel, ylabel, xlim, ylim, xscale, yscale, legend)\n        self.X, self.Y, self.fmts = None, None, fmts\n\n    def add(self, x, y):\n        # Add multiple data points into the figure\n        if not hasattr(y, \"__len__\"):\n            y = [y]\n        n = len(y)\n        if not hasattr(x, \"__len__\"):\n            x = [x] * n\n        if not self.X:\n            self.X = [[] for _ in range(n)]\n        if not self.Y:\n            self.Y = [[] for _ in range(n)]\n        for i, (a, b) in enumerate(zip(x, y)):\n            if a is not None and b is not None:\n                self.X[i].append(a)\n                self.Y[i].append(b)\n        self.axes[0].cla()\n        for x, y, fmt in zip(self.X, self.Y, self.fmts):\n            self.axes[0].plot(x, y, fmt)\n        self.config_axes()\n        display.display(self.fig)\n        display.clear_output(wait=True)\n        \nclass Timer: #@save\n    \"\"\"记录多次运行时间\"\"\"\n    def __init__(self):\n        self.times = []\n        self.start()\n    \n    def start(self):\n        \"\"\"启动计时器\"\"\"\n        self.tik = time.time()\n    \n    def stop(self):\n        \"\"\"停止计时器并将时间记录在列表中\"\"\"\n        self.times.append(time.time() - self.tik)\n        return self.times[-1]\n    \n    def avg(self):\n        \"\"\"返回平均时间\"\"\"\n        return sum(self.times) / len(self.times)\n    def sum(self):\n        \"\"\"返回时间总和\"\"\"\n        return sum(self.times)\n    \n    def cumsum(self):\n        \"\"\"返回累计时间\"\"\"\n        return np.array(self.times).cumsum().tolist()","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:29:15.874458Z","iopub.execute_input":"2024-08-18T08:29:15.874879Z","iopub.status.idle":"2024-08-18T08:29:15.903137Z","shell.execute_reply.started":"2024-08-18T08:29:15.874844Z","shell.execute_reply":"2024-08-18T08:29:15.902032Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 准备数据集\n### Preparing and spliting the data","metadata":{}},{"cell_type":"code","source":"#@save\nDATA_HUB = dict()\nDATA_URL = 'http://d2l-data.s3-accelerate.amazonaws.com/'\n\ndef download(name, cache_dir=os.path.join('..', 'data')): #@save\n    \"\"\"下载一个DATA_HUB中的文件，返回本地文件名\"\"\"\n    assert name in DATA_HUB, f\"{name} 不存在于 {DATA_HUB}\"\n    url, sha1_hash = DATA_HUB[name]\n    os.makedirs(cache_dir, exist_ok=True)\n    fname = os.path.join(cache_dir, url.split('/')[-1])\n    if os.path.exists(fname):\n        sha1 = hashlib.sha1()\n        with open(fname, 'rb') as f:\n            while True:\n                data = f.read(1048576)\n                if not data:\n                    break\n                sha1.update(data)\n        if sha1.hexdigest() == sha1_hash:\n            return fname # 命中缓存\n    print(f'正在从{url}下载{fname}...')\n    r = requests.get(url, stream=True, verify=True)\n    with open(fname, 'wb') as f:\n        f.write(r.content)\n    return fname\n\ndef download_extract(name, folder=None):\n    \"\"\"Download and extract a zip/tar file.\"\"\"\n    fname = download(name)\n    base_dir = os.path.dirname(fname)\n    data_dir, ext = os.path.splitext(fname)\n    if ext == '.zip':\n        fp = zipfile.ZipFile(fname, 'r')\n    elif ext in ('.tar', '.gz'):\n        fp = tarfile.open(fname, 'r')\n    else:\n        assert False, 'Only zip/tar files can be extracted.'\n    fp.extractall(base_dir)\n    return os.path.join(base_dir, folder) if folder else data_dir\n\nDATA_HUB['kaggle_dog_tiny'] = (DATA_URL + 'kaggle_dog_tiny.zip',\n                            '0cb91d09b814ecdc07b50f31f8dcad3e81d6a86d')\n# 如果使用完整的Kaggle竞赛的数据集，设置demo为False\ndemo = False\nif demo:\n    data_dir = download_extract('cifar10_tiny')\nelse:\n    data_dir = r'/kaggle/input/dog-breed-identification'","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:29:20.610089Z","iopub.execute_input":"2024-08-18T08:29:20.610732Z","iopub.status.idle":"2024-08-18T08:29:20.622373Z","shell.execute_reply.started":"2024-08-18T08:29:20.610699Z","shell.execute_reply":"2024-08-18T08:29:20.621409Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 整理数据集\n#@save\ndef read_csv_labels(fname):\n    \"\"\"读取fname来给标签字典返回一个文件名\"\"\"\n    with open(fname, 'r') as f:\n        # 跳过文件头行(列名)\n        lines = f.readlines()[1:]\n    tokens = [l.rstrip().split(',') for l in lines]\n    return dict(((name, label) for name, label in tokens))","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:29:27.434211Z","iopub.execute_input":"2024-08-18T08:29:27.434843Z","iopub.status.idle":"2024-08-18T08:29:27.440345Z","shell.execute_reply.started":"2024-08-18T08:29:27.434808Z","shell.execute_reply":"2024-08-18T08:29:27.439279Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#@save\ndef copyfile(filename, target_dir):\n    \"\"\"将文件复制到目标目录\"\"\"\n    os.makedirs(target_dir, exist_ok=True)\n    shutil.copy(filename, target_dir)\n    \n\"\"\"令n等于样本最少的类别\n中的图像数量，而r是比率。验证集将为每个类别拆分出max(⌊nr⌋, 1)张图像\n\"\"\"\n\n#@save\ndef reorg_train_valid(data_dir, labels, valid_ratio, target_dir):\n    \"\"\"将验证集从原始的训练集中拆分出来\"\"\"\n    # 训练数据集中样本最少的类别中的样本数\n    n = collections.Counter(labels.values()).most_common()[-1][1]\n    # 验证集中每个类别的样本数\n    n_valid_per_label = max(1, math.floor(n * valid_ratio))\n    label_count = {}\n    for train_file in os.listdir(os.path.join(data_dir, 'train')):\n        label = labels[train_file.split('.')[0]]\n        fname = os.path.join(data_dir, 'train', train_file)\n        copyfile(fname, os.path.join(target_dir, 'train_valid_test',\n                'train_valid', label))\n        if label not in label_count or label_count[label] < n_valid_per_label:\n            copyfile(fname, os.path.join(target_dir, 'train_valid_test',\n                    'valid', label))\n            label_count[label] = label_count.get(label, 0) + 1\n        else:\n            copyfile(fname, os.path.join(target_dir, 'train_valid_test',\n                    'train', label))\n    return n_valid_per_label\n\n#@save\ndef reorg_test(data_dir, target_dir):\n    \"\"\"在预测期间整理测试集，以方便读取\"\"\"\n    for test_file in os.listdir(os.path.join(data_dir, 'test')):\n        copyfile(os.path.join(data_dir, 'test', test_file),\n                os.path.join(target_dir, 'train_valid_test', 'test',\n                'unknown'))\n\n# 加载数据集分类标签\ndef reorg_cifar10_data(data_dir, valid_ratio,target_dir):\n    labels = read_csv_labels(os.path.join(data_dir, 'labels.csv'))\n    print('# 训练样本 :', len(labels))\n    print('# 类别 :', len(set(labels.values())))\n    reorg_train_valid(data_dir, labels, valid_ratio, target_dir)\n    reorg_test(data_dir, target_dir)","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:29:29.704642Z","iopub.execute_input":"2024-08-18T08:29:29.705481Z","iopub.status.idle":"2024-08-18T08:29:29.717200Z","shell.execute_reply.started":"2024-08-18T08:29:29.705446Z","shell.execute_reply":"2024-08-18T08:29:29.716211Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 初始化数据\nbatch_size = 32 if demo else 128\nvalid_ratio = 0.1\ntarget_dir = r'/kaggle/working/Dog-Breed-Identification'\nreorg_cifar10_data(data_dir, valid_ratio, target_dir)","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:29:33.216872Z","iopub.execute_input":"2024-08-18T08:29:33.217829Z","iopub.status.idle":"2024-08-18T08:30:19.555354Z","shell.execute_reply.started":"2024-08-18T08:29:33.217794Z","shell.execute_reply":"2024-08-18T08:30:19.554395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 图像增广\n### Image Augmentation","metadata":{}},{"cell_type":"code","source":"\"\"\"使用图像增广来解决过拟合的问题\"\"\"\n\"\"\"在训练期间，我们通常会对训练图像进行增广，以增加模型的准确性。\"\"\"\n\ntransform_train = torchvision.transforms.Compose([\n    # 随机裁剪图像，所得图像为原始面积的0.08～1之间，高宽比在3/4和4/3之间。\n    # 然后，缩放图像以创建224x224的新图像\n    torchvision.transforms.RandomResizedCrop(224, scale=(0.08, 1.0),\n                                            ratio=(3.0/4.0, 4.0/3.0)),\n    torchvision.transforms.RandomHorizontalFlip(),\n    # 随机更改亮度，对比度和饱和度\n    torchvision.transforms.ColorJitter(brightness=0.4,\n                                    contrast=0.4,\n                                    saturation=0.4),\n    torchvision.transforms.RandomAffine(degrees=(-30, 30), \n                                        translate=(0.05, 0.1), \n                                        shear=10),\n    # 添加随机噪声\n    torchvision.transforms.ToTensor(),\n    # 标准化图像的每个通道\n    torchvision.transforms.Normalize([0.485, 0.456, 0.406],\n                                    [0.229, 0.224, 0.225])])\n\ntransform_test = torchvision.transforms.Compose([\n    torchvision.transforms.Resize(256),\n    # 从图像中心裁切224x224大小的图片\n    torchvision.transforms.CenterCrop(224),\n    torchvision.transforms.ToTensor(),\n    torchvision.transforms.Normalize([0.485, 0.456, 0.406],\n                                    [0.229, 0.224, 0.225])])\n\n","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:30:27.551950Z","iopub.execute_input":"2024-08-18T08:30:27.552342Z","iopub.status.idle":"2024-08-18T08:30:27.560743Z","shell.execute_reply.started":"2024-08-18T08:30:27.552313Z","shell.execute_reply":"2024-08-18T08:30:27.559725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 读取数据同时进行增广\n\"\"\"在训练期间，我们需要指定上面定义的所有图像增广操作。当验证集在超参数调整过程中用于模型评估时，\n不应引入图像增广的随机性。\"\"\"\ntrain_ds, train_valid_ds = [torchvision.datasets.ImageFolder(\n                            os.path.join(target_dir, 'train_valid_test', folder),\n                            transform=transform_train) for folder in ['train', 'train_valid']]\nvalid_ds, test_ds = [torchvision.datasets.ImageFolder(\n                            os.path.join(target_dir, 'train_valid_test', folder),\n                            transform=transform_test) for folder in ['valid', 'test']]\n\n\"\"\"但是在最终预测之前，我们根据训练集和验证集组合而成的训练模型进行训练，以\n充分利用所有标记的数据。\n\"\"\"\ntrain_iter, train_valid_iter = [torch.utils.data.DataLoader(\n                                dataset, batch_size, shuffle=True, drop_last=True)\n                                for dataset in (train_ds, train_valid_ds)]\nvalid_iter = torch.utils.data.DataLoader(valid_ds, batch_size, shuffle=False,\n                                        drop_last=True)\ntest_iter = torch.utils.data.DataLoader(test_ds, batch_size, shuffle=False,\n                                        drop_last=False)","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:30:29.400580Z","iopub.execute_input":"2024-08-18T08:30:29.401447Z","iopub.status.idle":"2024-08-18T08:30:29.567562Z","shell.execute_reply.started":"2024-08-18T08:30:29.401411Z","shell.execute_reply":"2024-08-18T08:30:29.566707Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 构建模型\n### Build the Swin Transformer","metadata":{}},{"cell_type":"code","source":"# Swin from scratch\n# 原文链接：https://blog.csdn.net/guoqingru0311/article/details/128963854\n# Reference: https://blog.csdn.net/guoqingru0311/article/details/128963854\n\n\ndef drop_path_f(x, drop_prob: float = 0., training: bool = False):\n    \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\n\n    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,\n    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...\n    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for\n    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use\n    'survival rate' as the argument.\n\n    \"\"\"\n    if drop_prob == 0. or not training:\n        return x\n    keep_prob = 1 - drop_prob\n    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets\n    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)\n    random_tensor.floor_()  # binarize\n    output = x.div(keep_prob) * random_tensor\n    return output\n\n\nclass DropPath(nn.Module):\n    \"\"\"Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).\n    \"\"\"\n    def __init__(self, drop_prob=None):\n        super(DropPath, self).__init__()\n        self.drop_prob = drop_prob\n\n    def forward(self, x):\n        return drop_path_f(x, self.drop_prob, self.training)\n\n\"\"\"\n    将窗口按照window_size的大小进行划分，得到一个个窗口\n\"\"\"\ndef window_partition(x, window_size: int):\n    \"\"\"\n    将feature map按照window_size划分成一个个没有重叠的window\n    Args:\n        x: (B, H, W, C)\n        window_size (int): window size(M)\n\n    Returns:\n        windows: (num_windows*B, window_size, window_size, C)\n    \"\"\"\n    B, H, W, C = x.shape\n    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)\n    # permute: [B, H//Mh, Mh, W//Mw, Mw, C] -> [B, H//Mh, W//Mh, Mw, Mw, C]\n    # view: [B, H//Mh, W//Mw, Mh, Mw, C] -> [B*num_windows, Mh, Mw, C]\n    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)   # 输出的是按照指定的window_size划分成一个一个窗口的数据\n    return windows\n\n\ndef window_reverse(windows, window_size: int, H: int, W: int):\n    \"\"\"\n    将一个个window还原成一个feature map\n    Args:\n        windows: (num_windows*B, window_size, window_size, C)\n        window_size (int): Window size(M)\n        H (int): Height of image\n        W (int): Width of image\n\n    Returns:\n        x: (B, H, W, C)\n    \"\"\"\n    B = int(windows.shape[0] / (H * W / window_size / window_size))\n    # view: [B*num_windows, Mh, Mw, C] -> [B, H//Mh, W//Mw, Mh, Mw, C]\n    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)\n    # permute: [B, H//Mh, W//Mw, Mh, Mw, C] -> [B, H//Mh, Mh, W//Mw, Mw, C]\n    # view: [B, H//Mh, Mh, W//Mw, Mw, C] -> [B, H, W, C]\n    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)\n    return x\n\n\nclass PatchEmbed(nn.Module):\n    \"\"\"\n    2D Image to Patch Embedding\n    split image into non-overlapping patches   即将图片划分成一个个没有重叠的patch\n    \"\"\"\n    def __init__(self, patch_size=4, in_c=3, embed_dim=96, norm_layer=None):\n        super().__init__()\n        patch_size = (patch_size, patch_size)\n        self.patch_size = patch_size\n        self.in_chans = in_c\n        self.embed_dim = embed_dim\n        self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)\n        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()\n\n    def forward(self, x):\n        _, _, H, W = x.shape\n\n        # padding\n        # 如果输入图片的H，W不是patch_size的整数倍，需要进行padding\n        pad_input = (H % self.patch_size[0] != 0) or (W % self.patch_size[1] != 0)\n        if pad_input:\n            # to pad the last 3 dimensions,\n            # (W_left, W_right, H_top,H_bottom, C_front, C_back)\n            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1],   # 表示宽度方向右侧填充数\n                          0, self.patch_size[0] - H % self.patch_size[0],   # 表示高度方向底部填充数\n                          0, 0))\n\n        # 下采样patch_size倍\n        x = self.proj(x)\n        _, _, H, W = x.shape\n        # flatten: [B, C, H, W] -> [B, C, HW]\n        # transpose: [B, C, HW] -> [B, HW, C]\n        x = x.flatten(2).transpose(1, 2)\n        x = self.norm(x)\n        return x, H, W\n\n\nclass PatchMerging(nn.Module):\n    r\"\"\" Patch Merging Layer.\n        步长为2，间隔采样\n    Args:\n        dim (int): Number of input channels.\n        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n    \"\"\"\n\n    def __init__(self, dim, norm_layer=nn.LayerNorm):\n        super().__init__()\n        self.dim = dim\n        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)\n        self.norm = norm_layer(4 * dim)\n\n    def forward(self, x, H, W):\n        \"\"\"\n        x: B, H*W, C    即输入x的通道排列顺序\n        \"\"\"\n        B, L, C = x.shape\n        assert L == H * W, \"input feature has wrong size\"\n\n        x = x.view(B, H, W, C)\n\n        # padding\n        # 如果输入feature map的H，W不是2的整数倍，需要进行padding\n        pad_input = (H % 2 == 1) or (W % 2 == 1)\n        if pad_input:\n            # to pad the last 3 dimensions, starting from the last dimension and moving forward.\n            # (C_front, C_back, W_left, W_right, H_top, H_bottom)\n            # 注意这里的Tensor通道是[B, H, W, C]，所以会和官方文档有些不同\n            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))\n\n        # 以2为间隔进行采样\n        x0 = x[:, 0::2, 0::2, :]  # [B, H/2, W/2, C]\n        x1 = x[:, 1::2, 0::2, :]  # [B, H/2, W/2, C]\n        x2 = x[:, 0::2, 1::2, :]  # [B, H/2, W/2, C]\n        x3 = x[:, 1::2, 1::2, :]  # [B, H/2, W/2, C]\n        x = torch.cat([x0, x1, x2, x3], -1)  #  ————————>  [B, H/2, W/2, 4*C]   在channael维度上进行拼接\n        x = x.view(B, -1, 4 * C)  # [B, H/2*W/2, 4*C]\n\n        x = self.norm(x)\n        x = self.reduction(x)  # [B, H/2*W/2, 2*C]\n\n        return x\n\n\"\"\"\nMLP模块\n\"\"\"\nclass Mlp(nn.Module):\n    \"\"\" MLP as used in Vision Transformer, MLP-Mixer and related networks\n    \"\"\"\n    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n        super().__init__()\n        out_features = out_features or in_features\n        hidden_features = hidden_features or in_features\n\n        self.fc1 = nn.Linear(in_features, hidden_features)\n        self.act = act_layer()\n        self.drop1 = nn.Dropout(drop)\n        self.fc2 = nn.Linear(hidden_features, out_features)\n        self.drop2 = nn.Dropout(drop)\n\n    def forward(self, x):\n        x = self.fc1(x)\n        x = self.act(x)\n        x = self.drop1(x)\n        x = self.fc2(x)\n        x = self.drop2(x)\n        return x\n\n\"\"\"\nWindowAttention即为：SW-MSA或者W-MSA模块\n\"\"\"\nclass WindowAttention(nn.Module):\n    r\"\"\" Window based multi-head self attention (W-MSA) module with relative position bias.\n    It supports both of shifted and non-shifted window.\n\n    Args:\n        dim (int): Number of input channels.\n        window_size (tuple[int]): The height and width of the window.\n        num_heads (int): Number of attention heads.\n        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True\n        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0\n        proj_drop (float, optional): Dropout ratio of output. Default: 0.0\n    \"\"\"\n\n    def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.):\n\n        super().__init__()\n        self.dim = dim\n        self.window_size = window_size  # [Mh, Mw]\n        self.num_heads = num_heads\n        head_dim = dim // num_heads\n        self.scale = head_dim ** -0.5\n\n        # define a parameter table of relative position bias\n        # 创建偏置bias项矩阵\n        self.relative_position_bias_table = nn.Parameter(\n            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # [2*Mh-1 * 2*Mw-1, nH]    其元素的个数===>>[(2*Mh-1) * (2*Mw-1)]\n\n        # get pair-wise relative position index for each token inside the window\n        coords_h = torch.arange(self.window_size[0])  # 如果此处的self.window_size[0]为2的话，则生成的coords_h为[0,1]\n        coords_w = torch.arange(self.window_size[1])  # 同理得\n        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # [2, Mh, Mw]\n        coords_flatten = torch.flatten(coords, 1)  # [2, Mh*Mw]\n        # [2, Mh*Mw, 1] - [2, 1, Mh*Mw]\n        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # [2, Mh*Mw, Mh*Mw]\n        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # [Mh*Mw, Mh*Mw, 2]\n        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0  行标+（M-1）\n        relative_coords[:, :, 1] += self.window_size[1] - 1     # 列表标+（M-1）\n        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1\n        relative_position_index = relative_coords.sum(-1)  # [Mh*Mw, Mh*Mw]\n        self.register_buffer(\"relative_position_index\", relative_position_index)   # 将relative_position_index放入到模型的缓存当中\n\n        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n        self.attn_drop = nn.Dropout(attn_drop)\n        self.proj = nn.Linear(dim, dim)\n        self.proj_drop = nn.Dropout(proj_drop)\n\n        nn.init.trunc_normal_(self.relative_position_bias_table, std=.02)\n        self.softmax = nn.Softmax(dim=-1)\n\n    def forward(self, x, mask: Optional[torch.Tensor] = None):\n        \"\"\"\n        Args:\n            x: input features with shape of (num_windows*B, Mh*Mw, C)\n            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None\n        \"\"\"\n        # [batch_size*num_windows, Mh*Mw, total_embed_dim]\n        B_, N, C = x.shape\n        # qkv(): -> [batch_size*num_windows, Mh*Mw, 3 * total_embed_dim]\n        # reshape: -> [batch_size*num_windows, Mh*Mw, 3, num_heads, embed_dim_per_head]\n        # permute: -> [3, batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]\n        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n        # [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]\n        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)\n\n        # transpose: -> [batch_size*num_windows, num_heads, embed_dim_per_head, Mh*Mw]\n        # @: multiply -> [batch_size*num_windows, num_heads, Mh*Mw, Mh*Mw]\n        q = q * self.scale\n        attn = (q @ k.transpose(-2, -1))\n\n        # relative_position_bias_table.view: [Mh*Mw*Mh*Mw,nH] -> [Mh*Mw,Mh*Mw,nH]\n        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(\n            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)\n        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # [nH, Mh*Mw, Mh*Mw]\n        attn = attn + relative_position_bias.unsqueeze(0)\n\n        # 进行mask，相同区域使用0表示；不同区域使用-100表示\n        if mask is not None:\n            # mask: [nW, Mh*Mw, Mh*Mw]\n            nW = mask.shape[0]  # num_windows\n            # attn.view: [batch_size, num_windows, num_heads, Mh*Mw, Mh*Mw]\n            # mask.unsqueeze: [1, nW, 1, Mh*Mw, Mh*Mw]\n            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)\n            attn = attn.view(-1, self.num_heads, N, N)\n            attn = self.softmax(attn)\n        else:\n            attn = self.softmax(attn)\n\n        attn = self.attn_drop(attn)\n\n        # @: multiply -> [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]\n        # transpose: -> [batch_size*num_windows, Mh*Mw, num_heads, embed_dim_per_head]\n        # reshape: -> [batch_size*num_windows, Mh*Mw, total_embed_dim]\n        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)\n        x = self.proj(x)\n        x = self.proj_drop(x)\n        return x\n\n\"\"\"\n    SwinTransformerBlock\n\"\"\"\nclass SwinTransformerBlock(nn.Module):\n    r\"\"\" Swin Transformer Block.\n\n    Args:\n        dim (int): Number of input channels.\n        num_heads (int): Number of attention heads.\n        window_size (int): Window size.\n        shift_size (int): Shift size for SW-MSA.\n        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n        drop (float, optional): Dropout rate. Default: 0.0\n        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n        drop_path (float, optional): Stochastic depth rate. Default: 0.0\n        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU\n        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n    \"\"\"\n\n    def __init__(self, dim, num_heads, window_size=7, shift_size=0,\n                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0.,\n                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):\n        super().__init__()\n        self.dim = dim\n        self.num_heads = num_heads\n        self.window_size = window_size\n        self.shift_size = shift_size\n        self.mlp_ratio = mlp_ratio\n        assert 0 <= self.shift_size < self.window_size, \"shift_size must in 0-window_size\"\n\n        self.norm1 = norm_layer(dim)    # 先经过层归一化处理\n\n        # WindowAttention即为：SW-MSA或者W-MSA模块\n        self.attn = WindowAttention(\n            dim, window_size=(self.window_size, self.window_size), num_heads=num_heads, qkv_bias=qkv_bias,\n            attn_drop=attn_drop, proj_drop=drop)\n\n        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n        self.norm2 = norm_layer(dim)\n        mlp_hidden_dim = int(dim * mlp_ratio)\n        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\n    def forward(self, x, attn_mask):\n        H, W = self.H, self.W\n        B, L, C = x.shape\n        assert L == H * W, \"input feature has wrong size\"\n\n        shortcut = x\n        x = self.norm1(x)\n        x = x.view(B, H, W, C)\n\n        # pad feature maps to multiples of window size\n        # 把feature map给pad到window size的整数倍\n        pad_l = pad_t = 0\n        pad_r = (self.window_size - W % self.window_size) % self.window_size\n        pad_b = (self.window_size - H % self.window_size) % self.window_size\n        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))\n        _, Hp, Wp, _ = x.shape\n\n        # cyclic shift\n        # 判断是进行SW-MSA或者是W-MSA模块\n        if self.shift_size > 0:\n            # https://blog.csdn.net/ooooocj/article/details/126046858?ops_request_misc=&request_id=&biz_id=102&utm_term=torch.roll()%E7%94%A8%E6%B3%95&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-0-126046858.142^v73^control,201^v4^add_ask,239^v1^control&spm=1018.2226.3001.4187\n            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))    #进行数据移动操作\n        else:\n            shifted_x = x\n            attn_mask = None\n\n        # partition windows\n        # 将窗口按照window_size的大小进行划分，得到一个个窗口\n        x_windows = window_partition(shifted_x, self.window_size)  # [nW*B, Mh, Mw, C]\n        # 将数据进行展平操作\n        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # [nW*B, Mh*Mw, C]\n\n        # W-MSA/SW-MSA\n        \"\"\"\n            # 进行多头自注意力机制操作\n        \"\"\"\n        attn_windows = self.attn(x_windows, mask=attn_mask)  # [nW*B, Mh*Mw, C]\n\n        # merge windows\n        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)  # [nW*B, Mh, Mw, C]\n        # 将多窗口拼接回大的featureMap\n        shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # [B, H', W', C]\n\n        # reverse cyclic shift\n        # 将移位的数据进行还原\n        if self.shift_size > 0:\n            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))\n        else:\n            x = shifted_x\n        # 如果进行了padding操作，需要移出掉相应的pad\n        if pad_r > 0 or pad_b > 0:\n            # 把前面pad的数据移除掉\n            x = x[:, :H, :W, :].contiguous()\n\n        x = x.view(B, H * W, C)\n\n        # FFN\n        x = shortcut + self.drop_path(x)\n        x = x + self.drop_path(self.mlp(self.norm2(x)))\n\n        return x\n\n\nclass BasicLayer(nn.Module):\n    \"\"\"\n    A basic Swin Transformer layer for one stage.\n\n    Args:\n        dim (int): Number of input channels.\n        depth (int): Number of blocks.\n        num_heads (int): Number of attention heads.\n        window_size (int): Local window size.\n        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n        drop (float, optional): Dropout rate. Default: 0.0\n        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.\n    \"\"\"\n\n    def __init__(self, dim, depth, num_heads, window_size,\n                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0.,\n                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):\n        super().__init__()\n        self.dim = dim\n        self.depth = depth\n        self.window_size = window_size\n        self.use_checkpoint = use_checkpoint\n        self.shift_size = window_size // 2  # 表示向右和向下偏移的窗口大小   即窗口大小除以2，然后向下取整\n\n        # build blocks\n        self.blocks = nn.ModuleList([\n            SwinTransformerBlock(\n                dim=dim,\n                num_heads=num_heads,\n                window_size=window_size,\n                shift_size=0 if (i % 2 == 0) else self.shift_size,   # 通过判断shift_size是否等于0，来决定是使用W-MSA与SW-MSA\n                mlp_ratio=mlp_ratio,\n                qkv_bias=qkv_bias,\n                drop=drop,\n                attn_drop=attn_drop,\n                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,\n                norm_layer=norm_layer)\n            for i in range(depth)])\n\n        # patch merging layer    即：PatchMerging类\n        if downsample is not None:\n            self.downsample = downsample(dim=dim, norm_layer=norm_layer)\n        else:\n            self.downsample = None\n\n    def create_mask(self, x, H, W):\n        # calculate attention mask for SW-MSA\n        # 保证Hp和Wp是window_size的整数倍\n        Hp = int(np.ceil(H / self.window_size)) * self.window_size\n        Wp = int(np.ceil(W / self.window_size)) * self.window_size\n        # 拥有和feature map一样的通道排列顺序，方便后续window_partition\n        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # [1, Hp, Wp, 1]\n        h_slices = (slice(0, -self.window_size),\n                    slice(-self.window_size, -self.shift_size),\n                    slice(-self.shift_size, None))\n        w_slices = (slice(0, -self.window_size),\n                    slice(-self.window_size, -self.shift_size),\n                    slice(-self.shift_size, None))\n        cnt = 0\n        for h in h_slices:\n            for w in w_slices:\n                img_mask[:, h, w, :] = cnt\n                cnt += 1\n\n        # 将img_mask划分成一个一个窗口\n        mask_windows = window_partition(img_mask, self.window_size)  # [nW, Mh, Mw, 1]           # 输出的是按照指定的window_size划分成一个一个窗口的数据\n        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)  # [nW, Mh*Mw]\n        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)  # [nW, 1, Mh*Mw] - [nW, Mh*Mw, 1]  使用了广播机制\n        # [nW, Mh*Mw, Mh*Mw]\n        # 因为需要求得的是自身注意力机制，所以，所以相同的区域使用0表示，；不同的区域不等于0，填入-100，这样，在求得\n        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))   # 即对于不等于0的位置，赋值为-100；否则为0\n        return attn_mask\n\n    def forward(self, x, H, W):\n        attn_mask = self.create_mask(x, H, W)  # [nW, Mh*Mw, Mh*Mw]   # 制作mask蒙版\n        for blk in self.blocks:\n            blk.H, blk.W = H, W\n            if not torch.jit.is_scripting() and self.use_checkpoint:\n                x = checkpoint.checkpoint(blk, x, attn_mask)\n            else:\n                x = blk(x, attn_mask)\n        if self.downsample is not None:\n            x = self.downsample(x, H, W)\n            H, W = (H + 1) // 2, (W + 1) // 2\n\n        return x, H, W\n\n\nclass SwinTransformer(nn.Module):\n    r\"\"\" Swin Transformer\n        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -\n          https://arxiv.org/pdf/2103.14030\n\n    Args:\n        patch_size (int | tuple(int)): Patch size. Default: 4   表示通过Patch Partition层后，下采样几倍\n        in_chans (int): Number of input image channels. Default: 3\n        num_classes (int): Number of classes for classification head. Default: 1000\n        embed_dim (int): Patch embedding dimension. Default: 96\n        depths (tuple(int)): Depth of each Swin Transformer layer.\n        num_heads (tuple(int)): Number of attention heads in different layers.\n        window_size (int): Window size. Default: 7\n        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4\n        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True\n        drop_rate (float): Dropout rate. Default: 0\n        attn_drop_rate (float): Attention dropout rate. Default: 0\n        drop_path_rate (float): Stochastic depth rate. Default: 0.1\n        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.\n        patch_norm (bool): If True, add normalization after patch embedding. Default: True\n        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False\n    \"\"\"\n\n    def __init__(self, patch_size=4,  # 表示通过Patch Partition层后，下采样几倍\n                 in_chans=3,           # 输入图像通道\n                 num_classes=1000,     # 类别数\n                 embed_dim=96,         # Patch partition层后的LinearEmbedding层映射后的维度，之后的几层都是该数的整数倍  分别是 C、2C、4C、8C\n                 depths=(2, 2, 6, 2),  # 表示每一个Stage模块内，Swin Transformer Block重复的次数\n                 num_heads=(3, 6, 12, 24),  # 表示每一个Stage模块内，Swin Transformer Block中采用的Multi-Head self-Attention的head的个数\n                 window_size=7,         # 表示W-MSA与SW-MSA所采用的window的大小\n                 mlp_ratio=4.,          # 表示MLP模块中，第一个全连接层增大的倍数\n                 qkv_bias=True,\n                 drop_rate=0.,          # 对应的PatchEmbed层后面的\n                 attn_drop_rate=0.,     # 对应于Multi-Head self-Attention模块中对应的dropRate\n                 drop_path_rate=0.1,    # 对应于每一个Swin-Transformer模块中采用的DropRate   其是慢慢的递增的，从0增长到drop_path_rate\n                 norm_layer=nn.LayerNorm,\n                 patch_norm=True,\n                 use_checkpoint=False, **kwargs):\n        super().__init__()\n\n        self.num_classes = num_classes\n        self.num_layers = len(depths)  # depths:表示重复的Swin Transoformer Block模块的次数  表示每一个Stage模块内，Swin Transformer Block重复的次数\n        self.embed_dim = embed_dim\n        self.patch_norm = patch_norm\n        # stage4输出特征矩阵的channels\n        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))\n        self.mlp_ratio = mlp_ratio\n\n        # split image into non-overlapping patches   即将图片划分成一个个没有重叠的patch\n        self.patch_embed = PatchEmbed(\n            patch_size=patch_size, in_c=in_chans, embed_dim=embed_dim,\n            norm_layer=norm_layer if self.patch_norm else None)\n        self.pos_drop = nn.Dropout(p=drop_rate)   # PatchEmbed层后面的Dropout层\n\n        # stochastic depth\n        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule\n\n        # build layers\n        self.layers = nn.ModuleList()\n        for i_layer in range(self.num_layers):\n            # 注意这里构建的stage和论文图中有些差异\n            # 这里的stage不包含该stage的patch_merging层，包含的是下个stage的\n            layers = BasicLayer(dim=int(embed_dim * 2 ** i_layer),  # 传入特征矩阵的维度，即channel方向的深度\n                                depth=depths[i_layer],              # 表示当前stage中需要堆叠的多少Swin Transformer Block\n                                num_heads=num_heads[i_layer],       # 表示每一个Stage模块内，Swin Transformer Block中采用的Multi-Head self-Attention的head的个数\n                                window_size=window_size,            # 表示W-MSA与SW-MSA所采用的window的大小\n                                mlp_ratio=self.mlp_ratio,           # 表示MLP模块中，第一个全连接层增大的倍数\n                                qkv_bias=qkv_bias,\n                                drop=drop_rate,                     # 对应的PatchEmbed层后面的\n                                attn_drop=attn_drop_rate,           # 对应于Multi-Head self-Attention模块中对应的dropRate\n                                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],     # 对应于每一个Swin-Transformer模块中采用的DropRate   其是慢慢的递增的，从0增长到drop_path_rate\n                                norm_layer=norm_layer,\n                                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,   # 判断是否是第四个，因为第四个Stage是没有PatchMerging层的\n                                use_checkpoint=use_checkpoint)\n            self.layers.append(layers)\n\n        self.norm = norm_layer(self.num_features)\n        self.avgpool = nn.AdaptiveAvgPool1d(1)   # 自适应的全局平均池化\n        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()\n\n        self.apply(self._init_weights)\n\n    def _init_weights(self, m):\n        if isinstance(m, nn.Linear):\n            nn.init.trunc_normal_(m.weight, std=.02)\n            if isinstance(m, nn.Linear) and m.bias is not None:\n                nn.init.constant_(m.bias, 0)\n        elif isinstance(m, nn.LayerNorm):\n            nn.init.constant_(m.bias, 0)\n            nn.init.constant_(m.weight, 1.0)\n\n    def forward(self, x):\n        # x: [B, L, C]\n        x, H, W = self.patch_embed(x)  # 对图像下采样4倍\n        x = self.pos_drop(x)\n\n        # 依次传入各个stage中\n        for layer in self.layers:\n            x, H, W = layer(x, H, W)\n\n        x = self.norm(x)  # [B, L, C]\n        x = self.avgpool(x.transpose(1, 2))  # [B, C, 1]\n        x = torch.flatten(x, 1)\n        x = self.head(x)   # 经过全连接层，得到输出\n        return x","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:30:31.871666Z","iopub.execute_input":"2024-08-18T08:30:31.872050Z","iopub.status.idle":"2024-08-18T08:30:31.967177Z","shell.execute_reply.started":"2024-08-18T08:30:31.872015Z","shell.execute_reply":"2024-08-18T08:30:31.965885Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# the swin-tiny, swin-small, and swin base\ndef swin_small_patch4_window7_224(num_classes: int = 1000, **kwargs):\n    # trained ImageNet-1K\n    # https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth\n    model = SwinTransformer(in_chans=3,\n                            patch_size=4,\n                            window_size=7,\n                            embed_dim=96,\n                            depths=(2, 2, 18, 2),\n                            num_heads=(3, 6, 12, 24),\n                            num_classes=num_classes,\n                            **kwargs)\n    return model\n\ndef swin_tiny_patch4_window7_224(num_classes: int = 1000, **kwargs):\n    # trained ImageNet-1K\n    # https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth\n    model = SwinTransformer(in_chans=3,\n                            patch_size=4,\n                            window_size=7,\n                            embed_dim=96,\n                            depths=(2, 2, 6, 2),\n                            num_heads=(3, 6, 12, 24),\n                            num_classes=num_classes,\n                            **kwargs)\n    return model\n\ndef swin_base_patch4_window7_224(num_classes: int = 1000, **kwargs):\n    # trained ImageNet-1K\n    # https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224.pth\n    model = SwinTransformer(in_chans=3,\n                            patch_size=4,\n                            window_size=7,\n                            embed_dim=128,\n                            depths=(2, 2, 18, 2),\n                            num_heads=(4, 8, 16, 32),\n                            num_classes=num_classes,\n                            **kwargs)\n    return model","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:30:39.618797Z","iopub.execute_input":"2024-08-18T08:30:39.619515Z","iopub.status.idle":"2024-08-18T08:30:39.625768Z","shell.execute_reply.started":"2024-08-18T08:30:39.619480Z","shell.execute_reply":"2024-08-18T08:30:39.624785Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_swin(devices, pretrained=True):\n    \"\"\"Constructs a Swin-B model.\n    Args:\n        pretrained (bool): If True, returns a model pre-trained on ImageNet\n    \"\"\"\n    if not pretrained:\n        model = swin_tiny_patch4_window7_224(num_classes=120)\n        model = model.to(devices[0])\n        return model\n    else:\n        #这里提供torchvision提供的官方预训练模型，\n        model = swin_b(weights = Swin_B_Weights.IMAGENET1K_V1)\n        \n        # 与ResNet不同，直接修改swin的分类头，增加一段从1000类到120类的全连接层\n        # If you're using a ResNet, you could simply use code like model.output_new, \n        # but in Swin,the model will skip the new layers you added\n        \n        # tips: why input 1024? please check the official head by print(model.head)\n        # 为什么输入1024个通道，因为Swin-B原本的分类头是1024到1000，可以通过print(model.head)查看\n        model.head = nn.Sequential(nn.Linear(1024, 1000),\n                                    nn.ReLU(),\n                                    nn.Linear(1000, 256),\n                                    nn.ReLU(),\n                                    nn.Linear(256, 120))\n        # 初始化参数\n        # 使用He初始化初始化新添加的线性层\n        # Using Kaiming initialization because of the utilization of relu,while partly because that \n        # I am a huge fan of Kaiming He.\n        for m in model.head.modules():\n            if isinstance(m, nn.Linear):\n                torch.nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n                if m.bias is not None:\n                    torch.nn.init.constant_(m.bias, 0)\n        \n        # 将模型参数分配给用于计算的CPU或GPU\n        model = model.to(devices[0])\n        # 冻结特征层参数\n        # this loop give you the ability to unfreeze layers.\n        \n        unfreeze_layers = 0\n\n        # 遍历模型的特征层，从后向前解冻,例如Swin-S有7个特征层，每一个架构具体有几个层，可以查看上述手搓Swin的过程，或者打印模型结构\n        for i, layer in enumerate(reversed(list(model.features.children()))):\n            if i < unfreeze_layers:\n                for param in layer.parameters():\n                    param.requires_grad = True\n            else:\n                for param in layer.parameters():\n                    param.requires_grad = False\n                # print(i)\n        \n        \n        return model\n\n\n# 交叉熵损失\nloss = nn.CrossEntropyLoss(reduction='none')\n\n# 仅对新增部分进行损失计算\ndef evaluate_loss(data_iter, net, devices):\n    l_sum, n = 0.0, 0\n    for features, labels in data_iter:\n        features, labels = features.to(devices[0]), labels.to(devices[0])\n        outputs = net(features)\n        l = loss(outputs, labels)\n        l_sum += l.sum()\n        n += labels.numel()\n    return (l_sum / n).to('cpu')\n","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:30:42.409616Z","iopub.execute_input":"2024-08-18T08:30:42.410035Z","iopub.status.idle":"2024-08-18T08:30:42.420583Z","shell.execute_reply.started":"2024-08-18T08:30:42.409992Z","shell.execute_reply":"2024-08-18T08:30:42.419541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 定义训练函数\n### Training","metadata":{}},{"cell_type":"code","source":"def accuracy(y_hat, y): #@save\n    \"\"\"计算预测正确的数量\"\"\"\n    if len(y_hat.shape) > 1 and y_hat.shape[1] > 1:\n        y_hat = y_hat.argmax(axis=1)\n    cmp = y_hat.type(y.dtype) == y\n    return float(cmp.type(y.dtype).sum())\n\n\ndef evaluate_accuracy_gpu(net, data_iter, device=None):\n    \"\"\"Compute the accuracy for a model on a dataset using a GPU.\"\"\"\n    if isinstance(net, nn.Module):\n        net.eval()  # Set the model to evaluation mode\n        if not device:\n            device = next(iter(net.parameters())).device\n            \n    # No. of correct predictions, no. of predictions\n    metric = Accumulator(2)\n    with torch.no_grad():\n        for X, y in data_iter:\n            if isinstance(X, list):\n                # Required for BERT Fine-tuning (to be covered later)\n                X = [x.to(device) for x in X]\n            else:\n                X = X.to(device)\n            y = y.to(device)\n            metric.add(accuracy(net(X), y), y.numel())\n    return metric[0] / metric[1]\n\ndef train_batch(net, X, y, loss, trainer, devices):\n    \"\"\"Train for a minibatch with mutiple GPUs (defined in Chapter 13).\"\"\"\n    if isinstance(X, list):\n        # Required for BERT fine-tuning (to be covered later)\n        X = [x.to(devices[0]) for x in X]\n    else:\n        X = X.to(devices[0])\n    y = y.to(devices[0])\n    net.train()\n    trainer.zero_grad()\n    pred = net(X)\n    l = loss(pred, y)\n    l.sum().backward()\n    trainer.step()\n    train_loss_sum = l.sum()\n    train_acc_sum = accuracy(pred, y)\n    return train_loss_sum, train_acc_sum\n\ndef try_all_gpus():\n    \"\"\"Return all available GPUs, or [cpu(),] if no GPU exists.\"\"\"\n    devices = [\n        torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]\n    return devices if devices else [torch.device('cpu')]","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:30:46.182277Z","iopub.execute_input":"2024-08-18T08:30:46.182654Z","iopub.status.idle":"2024-08-18T08:30:46.195649Z","shell.execute_reply.started":"2024-08-18T08:30:46.182621Z","shell.execute_reply":"2024-08-18T08:30:46.194674Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 检查新的分类头是否生效\n# check the model's new head\ndevices = try_all_gpus()\nmodel = get_swin(devices, pretrained=True) \nprint(model.head)\n# you can also check the whole structure of Swin Transformer\n# print(model)\n\n# check the output of the model,to see if the head is correctly modified\na = torch.rand(1,3,224,224)\na = a.to(devices[0])\nprint(model(a).shape)\n\n# check the feature layers' parameters\n#for param in model.parameters():\n#    print(param.requires_grad)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train(net, train_iter, valid_iter, num_epochs, lr, wd, devices, lr_period, lr_decay):\n    # 只训练小型自定义输出网络\n    net = nn.DataParallel(net, device_ids=devices).to(devices[0])\n    #trainer = torch.optim.SGD((param for param in net.parameters()\n    #                         if param.requires_grad), lr=lr,\n    #                         momentum=0.9, weight_decay=wd)\n    trainer = torch.optim.AdamW((param for param in net.parameters()\n                            if param.requires_grad), lr=lr, weight_decay=wd)\n    scheduler = torch.optim.lr_scheduler.StepLR(trainer, lr_period, lr_decay)\n\n    num_batches, timer = len(train_iter), Timer()\n    legend = ['train loss']\n    if valid_iter is not None:\n        legend.append('valid loss')\n    animator = Animator(xlabel='epoch', xlim=[1, num_epochs],\n                        legend=legend)\n        \n    for epoch in range(num_epochs):\n        metric = Accumulator(2)\n        for i, (features, labels) in enumerate(train_iter):\n            timer.start()\n            features, labels = features.to(devices[0]), labels.to(devices[0])\n            trainer.zero_grad()\n            output = net(features)\n            l = loss(output, labels).sum()\n            l.backward()\n            trainer.step()\n            metric.add(l, labels.shape[0])\n            timer.stop()\n            if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:\n                            animator.add(epoch + (i + 1) / num_batches,\n                            (metric[0] / metric[1], None))\n        measures = f'train loss {metric[0] / metric[1]:.3f}'\n        if valid_iter is not None:\n            valid_loss = evaluate_loss(valid_iter, net, devices)\n            animator.add(epoch + 1, (None, valid_loss.detach().cpu()))\n        scheduler.step()\n    if valid_iter is not None:\n        measures += f', valid loss {valid_loss:.3f}'\n    print(measures + f'\\n{metric[1] * num_epochs / timer.sum():.1f}'\n          f' examples/sec on {str(devices)}')","metadata":{"execution":{"iopub.status.busy":"2024-08-18T10:05:33.340710Z","iopub.execute_input":"2024-08-18T10:05:33.341125Z","iopub.status.idle":"2024-08-18T10:05:33.353879Z","shell.execute_reply.started":"2024-08-18T10:05:33.341081Z","shell.execute_reply":"2024-08-18T10:05:33.352842Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# first run with valid dataset\n# valid loss will be smaller than train loss because of dropout in the Swin Transformer\ndevices, num_epochs, lr, wd = try_all_gpus(), 15, 1e-4, 4e-4\nlr_period, lr_decay, net = 5, 0.9, get_swin(devices, pretrained=True)\n# I have tested the P100 and T4 on Kaggle to train entie Swin Transformer\n# Sadly, it seems like that the params are too huge to tarin on Kaggle\ntrain(net, train_iter, valid_iter, num_epochs, lr, wd, devices, lr_period,\nlr_decay)","metadata":{"execution":{"iopub.status.busy":"2024-08-18T08:30:52.204706Z","iopub.execute_input":"2024-08-18T08:30:52.205393Z","iopub.status.idle":"2024-08-18T09:54:33.501604Z","shell.execute_reply.started":"2024-08-18T08:30:52.205357Z","shell.execute_reply":"2024-08-18T09:54:33.500539Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"devices, num_epochs, lr, wd = try_all_gpus(), 20, 1e-4, 4e-4\nlr_period, lr_decay = 5, 0.9\nnet = get_swin(devices, pretrained=True)\ntrain(net, train_valid_iter, None, num_epochs, lr, wd, devices, lr_period,\n                                                                lr_decay)","metadata":{"execution":{"iopub.status.busy":"2024-08-18T10:05:36.053540Z","iopub.execute_input":"2024-08-18T10:05:36.054198Z","iopub.status.idle":"2024-08-18T11:30:13.295720Z","shell.execute_reply.started":"2024-08-18T10:05:36.054163Z","shell.execute_reply":"2024-08-18T11:30:13.294414Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 保存参数\ntorch.save(net.state_dict(), '/kaggle/working/Swin-B_Dog.pth')","metadata":{"execution":{"iopub.status.busy":"2024-08-18T09:54:41.022830Z","iopub.execute_input":"2024-08-18T09:54:41.023297Z","iopub.status.idle":"2024-08-18T09:54:41.359048Z","shell.execute_reply.started":"2024-08-18T09:54:41.023256Z","shell.execute_reply":"2024-08-18T09:54:41.358020Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"preds = []\nfor data, label in test_iter:\n    output = torch.nn.functional.softmax(net(data.to(devices[0])), dim=1)\n    preds.extend(output.cpu().detach().numpy())\nids = sorted(os.listdir(\n            os.path.join(target_dir, 'train_valid_test', 'test', 'unknown')))\nwith open('/kaggle/working/submission.csv', 'w') as f:\n    f.write('id,' + ','.join(train_valid_ds.classes) + '\\n')\n    for i, output in zip(ids, preds):\n        f.write(i.split('.')[0] + ',' + ','.join(\n            [str(num) for num in output]) + '\\n') ","metadata":{"execution":{"iopub.status.busy":"2024-08-18T11:30:36.940411Z","iopub.execute_input":"2024-08-18T11:30:36.941118Z","iopub.status.idle":"2024-08-18T11:32:12.858883Z","shell.execute_reply.started":"2024-08-18T11:30:36.941086Z","shell.execute_reply":"2024-08-18T11:32:12.858039Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"csv_file = 'submission.csv'  # CSV文件的名称\nzip_file = 'submission.zip'   # 压缩文件的名称\nworking_dir = '/kaggle/working'\n\n# 确保CSV文件存在于当前目录\nif not os.path.isfile(csv_file):\n    raise FileNotFoundError(f\"The file {csv_file} does not exist.\")\n\n# 创建ZIP文件并添加CSV文件\nwith zipfile.ZipFile(os.path.join(working_dir, zip_file), 'w') as zipf:\n    zipf.write(os.path.join(working_dir, csv_file))\n\nprint(f\"{zip_file} has been created.\")","metadata":{},"execution_count":null,"outputs":[]}]}