1
2
sampler_x = RandomSampler(ds_x, replacement=True, num_samples=n_iters_per_epoch * batch_size)
batch_sampler_x = BatchSampler(sampler_x, batch_size, drop_last=True)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def load_data_train(L, dspth='./data'):
    """
    加载训练数据集并进行预处理

    参数：
    L (int): 标签数量
    dspth (str): 数据集路径

    返回：
    data_x, label_x, data_u, label_u: 已处理的训练数据和标签
    """
    # 构建训练数据文件列表
    datalist = [
        osp.join(dspth, 'cifar-10-batches-py', 'data_batch_{}'.format(i+1))
        for i in range(5)
    ]
    data, labels = [], []
    # 逐个读取数据文件并加载数据
    for data_batch in datalist:
        with open(data_batch, 'rb') as fr:
            entry = pickle.load(fr, encoding='latin1')
            lbs = entry['labels'] if 'labels' in entry.keys() else entry['fine_labels']
            data.append(entry['data'])
            labels.append(lbs)
    data = np.concatenate(data, axis=0)  # 合并数据
    labels = np.concatenate(labels, axis=0)  # 合并标签
    n_labels = L // 10  # 每个类别的标签数量
    data_x, label_x, data_u, label_u = [], [], [], []
    for i in range(10):  # 遍历每个类别
        indices = np.where(labels == i)[0]
        np.random.shuffle(indices)  # 随机打乱索引
        inds_x, inds_u = indices[:n_labels], indices[n_labels:]
        data_x += [
            data[i].reshape(3, 32, 32).transpose(1, 2, 0)
            for i in inds_x
        ]
        label_x += [labels[i] for i in inds_x]
        data_u += [
            data[i].reshape(3, 32, 32).transpose(1, 2, 0)
            for i in inds_u
        ]
        label_u += [labels[i] for i in inds_u]
    return data_x, label_x, data_u, label_u  # 返回处理后的数据和标签

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def mixmatch():
    model, Loss_x, Loss_u = set_model(device)

    n_iters_per_epoch = n_imgs_per_epoch // batchsize
    # 第一步：对数据进行增强
    dataloader_train_x, dataloader_train_x = LoadAndAugment_Val(
        batchsize, n_iters_per_epoch, L=n_labels, K=n_guesses, root=data_path
    )
    # 第二步：标签猜测 + Sharpening
    lb_guessor = GuessAndSharpening(model, T=temperature)
    # 第三步：MixUp,将有标签数据和无标签数据进行混合
    mixuper = MixUp(mixup_alpha)
    # 第四步：论文中参考Mean Teacher的EMA方法，我也模仿论文中的EMA方法，使用EMA方法进行参数更新
    ema = EMA(model, ema_alpha)  # 创建EMA对象
    ...
    ...

1
2
3
4
# 第一步：对数据进行增强
    dataloader_train_x, dataloader_train_x = LoadAndAugment_Val(
        batchsize, n_iters_per_epoch, L=n_labels, K=n_guesses, root=data_path
    )

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Cifar10Augment(Dataset):
    """
    CIFAR-10数据集的自定义Dataset类。
    """

    def __init__(self, data, labels, n_guesses=1, is_train=True):
        super(Cifar10Augment, self).__init__()
        self.data, self.labels = data, labels
        self.n_guesses = n_guesses
        assert len(self.data) == len(self.labels)
        assert self.n_guesses >= 1
        # 数据预处理和增强
        mean, std = (-0.0172, -0.0356, -0.1069), (0.4940, 0.4869, 0.5231) # [-1, 1]
        if is_train:
            self.trans = T.Compose([
                T.Resize((32, 32)),
                T.PadandRandomCrop(border=4, cropsize=(32, 32)),
                T.RandomHorizontalFlip(p=0.5),
                T.Normalize(mean, std),
                T.ToTensor(),
            ])
        else:
            self.trans = T.Compose([
                T.Resize((32, 32)),
                T.Normalize(mean, std),
                T.ToTensor(),
            ])
            
    def __getitem__(self, idx):
        im, lb = self.data[idx], self.labels[idx]
        ims = []
        # 对每个样本进行n_guesses次预处理和增强
        for _ in range(self.n_guesses):
            im_trans = self.trans(im)
            ims.append(im_trans)
        return ims, lb

1
2
3
4
5
6
7
8
9
10
11
12
13
# 创建有标签数据集和无标签数据集
dataset_x = Cifar10Augment(
    data=data_x,
    labels=label_x,
    n_guesses=1, #有标记只进行一次增强
    is_train=True
)
dataset_u = Cifar10Augment(
    data=data_u,
    labels=label_u,
    n_guesses=K,
    is_train=True
)

1
2
# 第二步：标签猜测 + Sharpening
lb_guessor = GuessAndSharpening(model, T=temperature)

1
2
3
4
5
6
7
8
9
10
11
...
# guess
all_probs = []  # 用于存储所有样本的类别概率
for im in ims: # ims (list of torch.Tensor): 输入的图像数据列表
    if use_gpu:
        im = im.cuda()  # 将图像数据移动到GPU上
    logits = self.guessor(im)  # 使用模型进行前向传播，得到logits
    probs = torch.softmax(logits, dim=1)  # 对logits进行softmax操作，得到预测类别概率
    all_probs.append(probs)  # 将类别概率添加到列表中
qb = sum(all_probs) / len(all_probs)  # 对所有样本的类别概率进行求和并取平均
...

1
2
3
4
5
...
# sharpen
lbs_tem = torch.pow(qb, 1. / self.T)  # 对概率进行温度调节
lbs = lbs_tem / torch.sum(lbs_tem, dim=1, keepdim=True)  # 对概率进行归一化，得到伪标签
...

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import torch

# 这个标签猜测器通过对无标签样本的预测结果进行温度调节，生成伪标签，用于无监督学习中的训练。

use_gpu = False  # 设置是否使用GPU

class GuessAndSharpening(object):
    def __init__(self, model, T):
        """
        初始化标签猜测器
        
        参数：
        - model (torch.nn.Module): 用于进行标签猜测的模型
        - T (float): softmax温度参数
        """
        self.T = T  # softmax温度参数
        self.guessor = model  # 用于进行标签猜测的模型

    @torch.no_grad()
    def __call__(self, model, ims):
        """
        进行标签猜测
        
        参数：
        - model (torch.nn.Module): 要进行标签猜测的模型
        - ims (list of torch.Tensor): 输入的图像数据列表
        
        返回：
        - lbs (torch.Tensor): 伪标签，形状为[batch_size, num_classes]
        """
        # guess
        org_state = {
            k: v.clone().detach()
            for k, v in model.state_dict().items()
        }  # 备份模型当前状态
        is_train = self.guessor.training  # 记录模型当前是否处于训练状态
        self.guessor.train()  # 将模型设置为训练模式
        all_probs = []  # 用于存储所有样本的类别概率
        for im in ims:
            if use_gpu:
                im = im.cuda()  # 将图像数据移动到GPU上
            logits = self.guessor(im)  # 使用模型进行前向传播，得到logits
            probs = torch.softmax(logits, dim=1)  # 对logits进行softmax操作，得到类别概率
            all_probs.append(probs)  # 将类别概率添加到列表中
        qb = sum(all_probs) / len(all_probs)  # 对所有样本的类别概率进行求和并取平均

        # sharpen
        lbs_tem = torch.pow(qb, 1. / self.T)  # 对概率进行温度调节
        lbs = lbs_tem / torch.sum(lbs_tem, dim=1, keepdim=True)  # 对概率进行归一化，得到伪标签
        self.guessor.load_state_dict(org_state)  # 恢复模型之前的状态
        if is_train:
            self.guessor.train()  # 如果之前处于训练状态，则重新设置为训练模式
        else:
            self.guessor.eval()  # 如果之前处于评估状态，则设置为评估模式
        return lbs.detach()  # 返回伪标签，使用detach()方法使其与计算图脱离，避免梯度传播

1
2
# 第三步：MixUp,将有标签数据和无标签数据进行混合
mixuper = MixUp(mixup_alpha)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import torch

use_gpu = False  # 设置是否使用GPU

class MixUp(object):
    def __init__(self, alpha):
        """
        - alpha (float): MixUp参数，控制混合程度
        """
        self.beta_generator = torch.distributions.beta.Beta(alpha, alpha)

    def __call__(self, ims, lbs):
        bs = ims.size(0)  # 获取batch_size
        if use_gpu:
            lam = self.beta_generator.sample([bs, 1, 1, 1]).cuda()  # 从Beta分布中采样生成混合比例lam
        else:
            lam = self.beta_generator.sample([bs, 1, 1, 1])
            
        lam = torch.where(lam > (1. - lam), lam, (1. - lam))  # 将lam限制在[0.5, 1.0]之间
        indices = torch.randperm(bs)  # 生成随机排列的索引
        ims = lam * ims + (1. - lam) * ims[indices]  # 进行MixUp操作，线性插值
        lam = lam.view(-1, 1)  # 将lam转换成[batch_size, 1]的形状
        lbs = lam * lbs + (1. - lam) * lbs[indices]  # 对标签进行MixUp操作，线性插值
        return ims.detach(), lbs.detach()  # 返回增强后的图像数据和标签数据，使用detach()方法使其与原数据脱离计算图，避免梯度传播

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
...
Loss_x = CrossEntropyLoss().to(device)  # 使用交叉熵损失函数
Loss_u = nn.MSELoss().to(device)        # 使用均方误差损失函数
...

optim.zero_grad()  # 梯度清零
logits = model(ims)  # 输入模型进行前向传播
logits_x = logits[:batchsize]  # 获取有标签数据的输出结果
lbs_x = lbs[:batchsize]         # 获取有标签数据的标签
logits_u = logits[batchsize:]   # 获取无标签数据的输出结果
preds_u = torch.softmax(logits_u, dim=1)  # 对无标签数据进行softmax操作
lbs_u = lbs[batchsize:]         # 获取无标签数据的标签
loss_x = Loss_x(logits_x, lbs_x)  # 计算有标签数据的损失
loss_u = Loss_u(preds_u, lbs_u)   # 计算无标签数据的损失
lam_u = lambda_u + lambda_u_once * it  # 计算无标签数据的损失权重
loss = loss_x + lam_u * loss_u         # 计算总损失
loss.backward()                        # 反向传播
optim.step()                           # 更新模型参数
do_weight_decay(model, wd)             # 权重衰减
ema.update_params()                    # 更新EMA参

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import torch
import torch.distributed as dist
'''
这个类实现了指数移动平均（EMA）方法，用于在训练过程中平滑模型参数。
它会记录模型参数的当前状态，并根据衰减率将其更新到影子副本中，
然后在训练过程中使用影子副本来计算损失和梯度，以减少参数更新的方差，从而提高模型的泛化能力。
'''

class EMA(object):
    def __init__(self, model, alpha):
        """
        初始化EMA对象
        
        参数：
        - model (torch.nn.Module): 要应用EMA的模型
        - alpha (float): EMA的衰减率，控制新旧参数的权重
        """
        self.step = 0  # 当前训练步数
        self.model = model  # 要应用EMA的模型
        self.alpha = alpha  # EMA的衰减率
        self.shadow = self.get_model_state()  # 创建模型参数的影子副本
        self.backup = {}  # 用于备份模型参数的字典
        self.param_keys = [k for k, _ in self.model.named_parameters()]  # 获取模型参数的键列表
        self.buffer_keys = [k for k, _ in self.model.named_buffers()]  # 获取模型缓冲区的键列表

    def update_params(self):
        """
        更新EMA的参数
        """
        decay = min(self.alpha, (self.step + 1) / (self.step + 10))  # 计算衰减因子
        state = self.model.state_dict()  # 获取模型的当前状态
        for name in self.param_keys:
            # 使用EMA公式更新参数的影子副本
            self.shadow[name].copy_(
                decay * self.shadow[name] + (1 - decay) * state[name]
            )
        self.step += 1  # 更新步数

    def update_buffer(self):
        """
        更新EMA的缓冲区
        """
        state = self.model.state_dict()  # 获取模型的当前状态
        for name in self.buffer_keys:
            # 使用当前状态更新缓冲区的影子副本
            self.shadow[name].copy_(state[name])

    def apply_shadow(self):
        """
        将EMA的参数应用到模型中
        """
        self.backup = self.get_model_state()  # 备份模型当前的参数状态
        self.model.load_state_dict(self.shadow)  # 将EMA的影子副本加载到模型中

    def restore(self):
        """
        恢复模型参数为备份状态
        """
        self.model.load_state_dict(self.backup)  # 将模型的参数恢复为备份状态

    def get_model_state(self):
        """
        获取模型当前状态的副本
        """
        return {
            k: v.clone().detach()  # 使用clone()方法创建参数的副本，并使用detach()方法使其与计算图脱离，避免梯度传播
            for k, v in self.model.state_dict().items()  # 遍历模型参数的字典，并创建副本
        }

1
2
3
4
5
...
logits_x, logits_u = logits_x_u[:n_x], logits_x_u[n_x:]
lbs_x = lbs_x.long()
loss_x = Loss_x(logits_x, lbs_x)  # 计算有标签数据的损失
...

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
class PadandRandomCrop(object):
    '''
    将输入张量期望为形状为(H, W, 3)
    '''
    def __init__(self, border=4, cropsize=(32, 32)):
        self.border = border
        self.cropsize = cropsize

    def __call__(self, im):
        # 添加边界
        borders = [(self.border, self.border), (self.border, self.border), (0, 0)]
        canvas = np.pad(im, borders, mode='reflect')
        H, W, C = canvas.shape
        h, w = self.cropsize
        # 计算裁剪范围
        dh, dw = max(0, H - h), max(0, W - w)
        sh, sw = np.random.randint(0, dh), np.random.randint(0, dw)
        # 随机裁剪
        out = canvas[sh:sh + h, sw:sw + w, :]
        return out

class RandomHorizontalFlip(object):
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, im):
        # 随机水平翻转
        if np.random.rand() < self.p:
            im = im[:, ::-1, :]
        return im

class Resize(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, im):
        # 调整大小
        im = cv2.resize(im, self.size)
        return im

class Normalize(object):
    '''
    输入为像素值范围在[0, 255]之间，通道顺序为'rgb'
    '''
    def __init__(self, mean, std):
        self.mean = np.array(mean, np.float32).reshape(1, 1, -1)
        self.std = np.array(std, np.float32).reshape(1, 1, -1)

    def __call__(self, im):
        if len(im.shape) == 4:
            mean, std = self.mean[None, ...], self.std[None, ...]
        elif len(im.shape) == 3:
            mean, std = self.mean, self.std
        # 将像素值归一化到[0, 1]
        im = im.astype(np.float32) / 255.
        # 像素值标准化
        im -= mean
        im /= std
        return im

class ToTensor(object):
    def __init__(self):
        pass

    def __call__(self, im):
        if len(im.shape) == 4:
            # 将NumPy数组转换为张量并转置通道
            return torch.from_numpy(im.transpose(0, 3, 1, 2))
        elif len(im.shape) == 3:
            # 将NumPy数组转换为张量并转置通道
            return torch.from_numpy(im.transpose(2, 0, 1))


class Compose(object):
    def __init__(self, ops):
        self.ops = ops

    def __call__(self, im):
        # 应用所有操作
        for op in self.ops:
            im = op(im)
        return im

1
2
3
4
5
6
7
8
9
10
11
def cutout_func(img, pad_size, replace=(0, 0, 0)):
    replace = np.array(replace, dtype=np.uint8)
    H, W = img.shape[0], img.shape[1]
    rh, rw = np.random.random(2)
    pad_size = pad_size // 2
    ch, cw = int(rh * H), int(rw * W)
    x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)
    y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)
    out = img.copy()
    out[x1:x2, y1:y2, :] = replace
    return out

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
# 图像增强函数字典
func_dict = {
    'Identity': identity_func,
    'AutoContrast': autocontrast_func,
    'Equalize': equalize_func,
    'Rotate': rotate_func,
    'Solarize': solarize_func,
    'Color': color_func,
    'Contrast': contrast_func,
    'Brightness': brightness_func,
    'Sharpness': sharpness_func,
    'ShearX': shear_x_func,
    'TranslateX': translate_x_func,
    'TranslateY': translate_y_func,
    'Posterize': posterize_func,
    'ShearY': shear_y_func,
}

# 默认参数字典
translate_const = 10
MAX_LEVEL = 10
replace_value = (128, 128, 128)
arg_dict = {
    'Identity': none_level_to_args,
    'AutoContrast': none_level_to_args,
    'Equalize': none_level_to_args,
    'Rotate': rotate_level_to_args(MAX_LEVEL, replace_value),
    'Solarize': solarize_level_to_args(MAX_LEVEL),
    'Color': enhance_level_to_args(MAX_LEVEL),
    'Contrast': enhance_level_to_args(MAX_LEVEL),
    'Brightness': enhance_level_to_args(MAX_LEVEL),
    'Sharpness': enhance_level_to_args(MAX_LEVEL),
    'ShearX': shear_level_to_args(MAX_LEVEL, replace_value),
    'TranslateX': translate_level_to_args(
        translate_const, MAX_LEVEL, replace_value
    ),
    'TranslateY': translate_level_to_args(
        translate_const, MAX_LEVEL, replace_value
    ),
    'Posterize': posterize_level_to_args(MAX_LEVEL),
    'ShearY': shear_level_to_args(MAX_LEVEL, replace_value),
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class RandomAugment(object):
    def __init__(self, N=2, M=10):
        self.N = N
        self.M = M

    # 随机选择N个增强操作
    def get_random_ops(self):
        sampled_ops = np.random.choice(list(func_dict.keys()), self.N)
        return [(op, 0.5, self.M) for op in sampled_ops]

    def __call__(self, img):
        ops = self.get_random_ops()
        for name, prob, level in ops:
            if np.random.random() > prob: continue
            args = arg_dict[name](level)
            img = func_dict[name](img, *args)
        img = cutout_func(img, 16, replace_value)
        return img

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class Cifar10Augment(Dataset):
    """
    自定义CIFAR-10数据集类
    """
    def __init__(self, data, labels, is_train=True):
        super(Cifar10Augment, self).__init__()
        self.data, self.labels = data, labels
        self.is_train = is_train
        assert len(self.data) == len(self.labels)  # 确保数据和标签数量相等
        mean, std = (0.4914, 0.4822, 0.4465), (0.2471, 0.2435, 0.2616)
        #  mean, std = (-0.0172, -0.0356, -0.1069), (0.4940, 0.4869, 0.5231) # [-1, 1]
        if is_train:
            self.trans_weak = T.Compose([
                T.Resize((32, 32)),
                T.PadandRandomCrop(border=4, cropsize=(32, 32)),
                T.RandomHorizontalFlip(p=0.5),
                T.Normalize(mean, std),
                T.ToTensor(),
            ])
            self.trans_strong = T.Compose([
                T.Resize((32, 32)),
                T.PadandRandomCrop(border=4, cropsize=(32, 32)),
                T.RandomHorizontalFlip(p=0.5),
                RandomAugment(2, 10), # 强增强
                T.Normalize(mean, std),
                T.ToTensor(),
            ])
        else:
            self.trans = T.Compose([
                T.Resize((32, 32)),
                T.Normalize(mean, std),
                T.ToTensor(),
            ])

1
lbs_u, valid_u = lb_guessor(model, ims_u_weak)  # 使用LabelGuessor预测未标记数据的标签

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
import torch

class LabelGuessor(object):
    def __init__(self, thresh):
        self.thresh = thresh  # 设置阈值，用于筛选高置信度的标签

    def __call__(self, model, ims):
        # 备份模型的当前状态
        org_state = {
            k: v.clone().detach()
            for k, v in model.state_dict().items()
        }
        # 记录模型当前的训练模式（训练或评估）
        is_train = model.training
        with torch.no_grad():  # 禁用梯度计算
            model.train()  # 切换模型到训练模式，以确保BatchNorm和Dropout层的行为一致
            all_probs = []  # 初始化一个空列表，用于存储所有的概率
            logits = model(ims)  # 前向传播，计算模型的输出
            probs = torch.softmax(logits, dim=1)  # 计算每个类别的概率
            scores, lbs = torch.max(probs, dim=1)  # 获取每个样本的最高概率和对应的标签
            
            idx = scores > self.thresh  # 根据阈值筛选高置信度的样本
            lbs = lbs[idx]  # 仅保留高置信度样本的标签

        # 恢复模型的原始状态
        model.load_state_dict(org_state)
        if is_train:
            model.train()  # 恢复训练模式
        else:
            model.eval()  # 恢复评估模式
        return lbs.detach(), idx  # 返回高置信度样本的标签和索引

1
2
3
4
5
6
logits_x_u = model(ims_x_u)
logits_x, logits_u = logits_x_u[:n_x], logits_x_u[n_x:]
lbs_x = lbs_x.long()
loss_x = Loss_x(logits_x, lbs_x)  # 计算有标签数据的损失
loss_u = Loss_u(logits_u, lbs_u)  # 计算无标签数据的损失

1
loss = loss_x + lambda_u * loss_u  # 总损失

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# 超参数设置
wresnet_k = 2
wresnet_n = 28
n_classes = 10
n_labeled = 250
n_epoches = 20
batchsize = 64
mu = 7
thr = 0.95
n_imgs_per_epoch = 64 * 1000
lam_u = 1.0
ema_alpha = 0.999
lr = 0.03
weight_decay = 5e-4
momentum = 0.9
seed = -1

sample_interval = 500  # 每500个迭代输出一次训练信息
data_path = '../../Mixmatch/data'
# 全局设置
use_gpu = True
device = torch.device('cuda' if use_gpu else 'cpu')  # 设置设备为GPU或CPU