提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

前言

1.可以去github直接下载这个项目,这样下载得到的是比较干净的版本，我把有注释的按照放在本文中，参考学习

项目地址

2.本文的一些python基础可在下面的链接中的第一部分ImageMasking找到
python基础

---

从ipynb文件入手

带注释的python文件

modules

mask.py

from requests import patch
from timm.models.vision_transformer import PatchEmbed, Block
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------

from functools import partial
import random
import math
import torch
import torch.nn as nn
import numpy as np
from timm.models.vision_transformer import PatchEmbed, Block

from modules.pos_embed import get_2d_sincos_pos_embed


class MaskingStrategy(nn.Module):
    """ Masked Autoencoder with VisionTransformer backbone
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3,
                 embed_dim=1024, choice='random'):
        super().__init__()

        # --------------------------------------------------------------------------
        # MAE encoder specifics
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        self.choice=choice
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim), requires_grad=False)  # fixed sin-cos embedding
        self.norm = nn.LayerNorm(embed_dim)
        self.initialize_weights()

    def initialize_weights(self):
        # initialization
        # initialize (and freeze) pos_embed by sin-cos embedding
        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.patch_embed.num_patches**.5), cls_token=True)
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))


        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
        w = self.patch_embed.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
        torch.nn.init.normal_(self.cls_token, std=.02)

        # initialize nn.Linear and nn.LayerNorm
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            # we use xavier_uniform following official JAX ViT:
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def patchify(self, imgs):
        """
        imgs: (N, 3, H, W)
        x: (N, L, patch_size**2 *3)
        """
        p = self.patch_embed.patch_size[0]
        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0

        h = w = imgs.shape[2] // p
        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
        x = torch.einsum('nchpwq->nhwpqc', x)
        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))
        return x

    def unpatchify(self, x):
        """
        x: (N, L, patch_size**2 *3)
        imgs: (N, 3, H, W)
        """
        p = self.patch_embed.patch_size[0]
        h = w = int(x.shape[1]**.5)
        assert h * w == x.shape[1]
        
        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
        return imgs

    def random_masking(self, x, mask_ratio):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """
        N, L, D = x.shape  # batch, length, dim
        len_keep = int(L * (1 - mask_ratio))
        
        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
        
        # sort noise for each sample
        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
        # The location of i-th (0-L) patch in ids_shuffle
        ids_restore = torch.argsort(ids_shuffle, dim=1)

        # keep the first subset
        ids_keep = ids_shuffle[:, :len_keep]
        # only keep first unmasked embeddings via indexing 
        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))

        # generate the binary mask: 0 is keep, 1 is remove
        mask = torch.ones([N, L], device=x.device)
        mask[:, :len_keep] = 0
        # unshuffle to get the binary mask
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return x_masked, mask, ids_restore

    def forward(self, x, mask_ratio):
        # embed patches
        x = self.patch_embed(x)

        # add pos embed w/o cls token
        x = x + self.pos_embed[:, 1:, :]

        # masking: length -> length * mask_ratio
        if self.choice=='random':
            x, mask, ids_restore = self.random_masking(x, mask_ratio)
        else:
            pass

        # append cls token
        cls_token = self.cls_token + self.pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        x = self.norm(x)

        return x, mask, ids_restore




class RandomMaskingStrategy:
    """ Masked Autoencoder with VisionTransformer backbone
    """
    def __init__(self, img_size=224, patch_size=16, device='cpu'):
        super().__init__()

        # --------------------------------------------------------------------------
        self.patch_size=patch_size
        self.img_size=img_size
        # 这里计算这副图片能够划分为多少个patch
        self.num_patches=int((img_size/patch_size)**2)
        self.device=torch.device(device)
        # self.device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')


    def patchify(self, imgs):
        """
        看到这里的N，就是刚刚那个unsqueeze(0)的作用，匹配上这个维度
        imgs: (N, 3, H, W)
        x: (N, L, patch_size**2 *3)
        """
        p = self.patch_size
        # assert 语句用于测试表达式是否为真。如果表达式为假，则会引发 AssertionError 并终止程序执行。
        # 这一句就要求我们输入的图片高和宽相等，我们这里是224*224，并且高可以整除patch_size
        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0

        h = w = imgs.shape[2] // p
        #读这一句，随机损失读者寿命,作用就是将图片成功划分为patch，对应了函数名-patchify
        '''
        imgs.shape[0]：保持批量大小不变，即 N。
        3：保持通道数不变
        h 和 p：将高度 H 拆分为 h 和 p 两个维度。
        w 和 p：将宽度 W 拆分为 w 和 p 两个维度。
        '''
        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
        '''
        PyTorch 的 einsum 函数来重新排列张量的维度。
        'nchpwq' 表示输入张量 x 的当前维度排列，即 (N, C, h, p, w, q)。
        'nhwpqc' 表示输出张量的目标维度排列，即 (N, h, w, p, q, C)
        N 代表批处理维度
        h 代表分块后的高度
        w 代表分块后的宽度
        p 和 q 代表每个块的高度和宽度
        C 代表通道数
        '''
        x = torch.einsum('nchpwq->nhwpqc', x)
        '''
        imgs.shape[0]：保持批处理维度不变，即 N。
        h * w：将原来的高有几块和宽有几块，相乘表示总的图像块数。
        p**2 * 3：将每个图像块的高度、宽度和通道数合并。
        '''
        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))
        return x

    def unpatchify(self, x):
        """
        这个将被划分开的小图片集 重新组合为 一张完整的图片，与上面的函数是反操作
        x: (N, L, patch_size**2 *3)
        imgs: (N, 3, H, W)
        """
        p = self.patch_size
        h = w = int(x.shape[1]**.5)
        assert h * w == x.shape[1]
        
        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
        return imgs

    def random_masking(self, batch_size, mask_ratio):
        """
        Perform per-sample random masking by per-sample shuffling.
        Per-sample shuffling is done by argsort random noise.
        x: [N, L, D], sequence
        """
        # N, L, D = x.shape  # batch, length, dim
        N=batch_size
        L=self.num_patches
        #int() 向下取整
        len_keep = int(L * (1 - mask_ratio))
        # 生成一个形状为 (N, L) 的张量，张量中的值在 [0, 1) 的区间内均匀分布
        noise = torch.rand(N, L, device=self.device)  # noise in [0, 1]
        
        # sort noise for each sample
        #dim=1是根据L维度，就是对被划分的小块patch的random值进行排序，默认升序，排队从矮到高
        # 下面这两句讲解可见python基础1.4
        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
        # The location of i-th (0-L) patch in ids_shuffle
        # 将 ids_shuffle 还原到原始顺序
        ids_restore = torch.argsort(ids_shuffle, dim=1)

        # keep the first subset
        ids_keep = ids_shuffle[:, :len_keep]
        # only keep first unmasked embeddings via indexing 
        # x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))

        # generate the binary mask: 0 is keep, 1 is remove
        mask = torch.ones([N, L], device=self.device)
        mask[:, :len_keep] = 0
        # unshuffle to get the binary mask
        # 33 取消洗牌以获得二进制掩码
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return mask, ids_restore

    def forward(self, x, mask_ratio):
        batch_size=x.shape[0]

        mask, ids_restore = self.random_masking(batch_size, mask_ratio)

        patch_size=16
        # visualize the mask
        # mask.detach() 会创建一个新张量，该张量与原始张量 mask 具有相同的数据，但不会再参与梯度计算。
        mask = mask.detach() #[batch_size, H*W]
        '''
        unsqueeze(-1):张量在最后一个维度（-1 表示最后一个维度）上增加一个维度。
        repeat(1, 1, patch_size**2 * 3) 对 unsqueeze(-1) 后的张量进行重复操作。具体来说：
        第一个参数 1 表示沿着第一个维度（即 N）不重复，保持不变。
        第二个参数 1 表示沿着第二个维度（即 H*W）不重复，保持不变。
        第三个参数 patch_size**2 * 3 表示沿着第三个维度重复 patch_size**2 * 3 次。
        这里 patch_size 是一个标量，用来表示一个图像块的大小，
        patch_size**2 * 3 表示每个像素点对应的信息复制了 patch_size**2 个通道
        
        '''
        pix_mask = mask.unsqueeze(-1).repeat(1, 1, patch_size**2 *3)  # [N, H*W, p*p*3]
        pix_mask = self.unpatchify(pix_mask)  # 1 is removing, 0 is keeping
        # 将 pix_mask 张量按照指定的顺序进行维度转置，
        # 并确保返回的张量不再保留梯度信息，并将其移动到 CPU 上进行后续处理或输出
        pix_mask = torch.einsum('nchw->nhwc', pix_mask).detach().cpu()

        x = torch.einsum('nchw->nhwc', x)

        # masked image 0 is keep, 1 is remove
        im_masked = x * (1 - pix_mask)
        # 要访问返回值的第0个维度得到被遮掩的图像im_masked ，
        return im_masked, mask, ids_restore




class BlockMasking:
    def __init__(
            self, img_size, patch_size=16, device='cpu'):
        self.device=torch.device(device)
        self.patch_size=patch_size
        self.height = int(img_size/patch_size)
        self.width = int(img_size/patch_size)
        self.num_patches = self.height * self.width
        self.num_masking_patches=None
        self.min_num_patches=None
        self.max_num_patches=None
        # max_aspect = max_aspect or 1 / min_aspect
        self.log_aspect_ratio = None
        
    # 该方法用于定义该类的实例的“官方”字符串表示形式，见python基础1.5
    # 可以直观的观察到初始化的数据
    def __repr__(self):
        repr_str = "Generator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
            self.height, self.width, self.min_num_patches, self.max_num_patches,
            self.num_masking_patches, self.log_aspect_ratio[0], self.log_aspect_ratio[1])
        return repr_str

    # def get_shape(self):
    #     return self.height, self.width
    
    def unpatchify(self, x):
        """
        x: (batch_size, H,W,768)
        """

        p = self.patch_size
        h=w= x.shape[1]
        # h = w = int(x.shape[1]**.5)
        # assert h * w == x.shape[1]
        
        x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
        x = torch.einsum('nhwpqc->nchpwq', x)
        imgs = x.reshape(shape=(x.shape[0], 3, h * p, h * p))
        return imgs

    def _mask(self, mask, max_mask_patches):
        delta = 0
        for attempt in range(10):
            #生成两个指定参数之间的随机浮点数，包括下限而不包括上限
            target_area = random.uniform(self.min_num_patches, max_mask_patches)
            # 星号 * 在函数调用中的作用是将一个可迭代对象（比如列表或元组）解包成单独的位置参数。
            # self.log_aspect_ratio 应该是一个包含两个元素的列表或元组，用来指定 random.uniform 函数的参数范围。
            # 在指定比例范围内随机生成一个比例，aspect_ratio是高宽比
            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
            #round() - 四舍五入，距离相等靠近偶数，h和w是高和宽方向遮掩几块
            h = int(round(math.sqrt(target_area * aspect_ratio)))
            w = int(round(math.sqrt(target_area / aspect_ratio)))
            if w < self.width and h < self.height:
                top = random.randint(0, self.height - h)
                left = random.randint(0, self.width - w)

                num_masked = mask[top: top + h, left: left + w].sum()
                # Overlap
                if 0 < h * w - num_masked <= max_mask_patches:
                    for i in range(top, top + h):
                        for j in range(left, left + w):
                            if mask[i, j] == 0:
                                mask[i, j] = 1
                                delta += 1

                if delta > 0:
                    break
        return delta

    def __call__(self,x, mask_ratio=0.4, min_num_patches=4, max_num_patches=None,
            min_aspect=0.3, max_aspect=None):
        
        batch_size=x.shape[0]
        self.num_masking_patches = self.num_patches*mask_ratio
        self.min_num_patches = min_num_patches
        self.max_num_patches = self.num_patches*mask_ratio if max_num_patches is None else max_num_patches

        max_aspect = max_aspect or 1 / min_aspect
        self.log_aspect_ratio = (math.log(min_aspect), math.log(max_aspect))

        mask = np.zeros(shape=(batch_size,self.height,self.width), dtype=np.int64)
        for i in range(batch_size):
            mask_count = 0
            while mask_count < self.num_masking_patches:
                max_mask_patches = self.num_masking_patches - mask_count
                max_mask_patches = min(max_mask_patches, self.max_num_patches)

                delta = self._mask(mask[i], max_mask_patches)
                if delta == 0:
                    break
                else:
                    mask_count += delta
        mask=torch.from_numpy(mask)



        pix_mask = mask.unsqueeze(-1).repeat(1, 1,1, self.patch_size**2 *3)# (1, H, W, p*p*3)
        pix_mask = self.unpatchify(pix_mask)  # 1 is removing, 0 is keeping
        pix_mask = torch.einsum('nchw->nhwc', pix_mask).detach().cpu()


        x = torch.einsum('nchw->nhwc', x)

        im_masked = x * (1 - pix_mask)
        return im_masked,mask

utils.py

import torchvision.transforms.functional as F
import random
import torch
from PIL import Image
import numpy as np
import os
import datetime
import logging
import matplotlib.pyplot as plt

def RandomSeed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False



def show_image(image, title=''):
    # image is [H, W, 3]
    assert image.shape[2] == 3
    plt.imshow(torch.clip(image * 255, 0, 255).int())
    plt.title(title, fontsize=16)
    plt.axis('off')
    return


class PairedTransform:
    def __init__(self, transform):
        self.transform = transform

    def __call__(self, img1, img2):
        # 保存当前随机种子
        torch_state = torch.get_rng_state()
        random_state = random.getstate()

        # 对第一张图片应用变换
        img1 = self.transform(img1)

        # 恢复随机种子，确保相同的变换应用到第二张图片
        torch.set_rng_state(torch_state)
        random.setstate(random_state)
        img2 = self.transform(img2)

        return img1, img2
    

#保持比例，缩放图像
class ResizeWithAspectRatio:
    #关于类的__init__和__call__的使用，基础1.2有介绍

    def __init__(self, size, interpolation=Image.BILINEAR):
        """
        Initialize the transform.

        Args:
        这里的size参数是要传入元组类型的目标尺寸，如(224，128)
        interpolation是所选择的插值方法
        - size (tuple): Desired output size (height, width).
        - interpolation (int): Interpolation method. Default is PIL.Image.BILINEAR.
        """
        self.size = size
        self.interpolation = interpolation

    def __call__(self, img):
        """
        Apply the transform to an image.

        Args:
        - img (PIL Image or Tensor): Image to be resized.

        Returns:
        - PIL Image or Tensor: Resized image.
        """
        #获取原始的图像的尺寸和目标图像的尺寸
        original_width, original_height = img.size
        target_height, target_width = self.size

        # Calculate new size to keep aspect ratio
        #计算原始图像的宽高比，最后生成的图像是宽高比不变的，尽可能接近目标尺寸
        aspect_ratio = original_width / original_height

        # 以宽高中更大的那一个尺寸为基准，另一个尺寸根据宽高比计算得到
        if aspect_ratio > 1:  # Wider image
            new_width = target_width
            new_height = int(target_width / aspect_ratio)
        else:  # Taller image or square
            new_height = target_height
            new_width = int(target_height * aspect_ratio)

        # Resize the image
        # 将图像重新调整大小
        resized_image = F.resize(img, (new_height, new_width), self.interpolation)

        # Calculate padding
        # 计算等比例且尽可能靠近目标大小后(大尺寸那一边完全满足)，剩下一小部分需要填充，采取以图像居中，四周填充的策略
        pad_left = (target_width - new_width) // 2
        pad_right = target_width - new_width - pad_left
        pad_top = (target_height - new_height) // 2
        pad_bottom = target_height - new_height - pad_top

        # Pad the resized image to the target size
        #resized_image：需要填充的图像,(pad_left, pad_top, pad_right, pad_bottom)：定义填充的大小。分别表示左、上、右、下的填充像素数。
        # padding_mode='constant'：填充模式。constant 表示用常数值填充。
        # fill=0：填充值。此处用 0 填充，即用黑色填充边缘
        padded_image = F.pad(resized_image, (pad_left, pad_top, pad_right, pad_bottom), padding_mode='constant', fill=0)

        return padded_image

def setup_logging(work_dir,name):
    # Create a directory with the current timestamp
    current_time = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
    run_dir = os.path.join(work_dir, f"{name}_{current_time}")
    os.makedirs(run_dir, exist_ok=True)

    # Set up logging to file and console
    log_file = os.path.join(run_dir, 'log.txt')
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[
        logging.FileHandler(log_file),
        logging.StreamHandler()
    ])

    return run_dir