cifar10_4000跑出论文精度代码

docs/zh_CN/training/semi_supervised_learning/FixMatchCCSSL.md

+# CCSSL:Class-Aware Contrastive Semi-Supervised Learning
+论文出处：[https://arxiv.org/abs/2203.02261](https://arxiv.org/abs/2203.02261)
+
+## 目录
+* 1. 原理介绍
+* 2. 精度指标
+* 3. 数据准备
+* 4. 模型训练
+* 5. 模型评估与推理部署
+* 5.1 模型评估
+* 5.2 模型推理
+* * 5.2.1 推理模型准备
+* * 5.2.2 基于Python预测引擎推理
+* * 5.2.3 基于C++预测引擎推理
+* 5.4 服务化部署
+* 5.5 端侧部署
+* 5.6 Paddle2ONNX模型转换与预测
+* 6. 参考文献
+
+## 1. 原理介绍
+作者提出了一种新颖的半监督学习方法。对有标签的数据进行数据训练的同时，对无标签数据进行一种弱增强和两种强增强。如果若增强的分类结果大于阈值，则弱数据增强的输出标签作为伪标签。通过伪标签，制作一个仅包含类级信息的监督对比矩阵。然后，通过对分布外数据的图像级对比形成类感知对比矩阵，以减少确认偏差。通过应用重新加权模块，将学习重点放在干净的数据上，并获得最终的目标矩阵。此外，特征亲和矩阵由两个强大的增强视图组成。通过最小化亲和矩阵和目标矩阵之间的交叉熵来制定用于未标记数据的类感知对比模块。模型的流程图如下
+
+## 2. 精度指标
+以下表格总结了复现的CCSSL在Cifar10数据集上的精度指标。
+<table>
+    <tr>
+        <td>Labels</td>
+        <td>40</td>
+        <td>250</td>
+        <td>4000</td>
+    </tr>
+    <tr>
+        <td>pytorch版本</td>
+        <td></td>
+        <td></td>
+        <td></td>
+    </tr>
+    <tr>
+        <td>paddle版本</td>
+        <td></td>
+        <td></td>
+        <td></td>
+    </tr>
+</table>
+cifar10上，paddle版本的配置文件及训练好的模型如下表所示
+<table>
+    <tr>
+        <td>label</td>
+        <td>配置文件地址</td>
+        <td>模型下载链接</td>
+    </tr>
+    <tr>
+        <td>40</td>
+        <td></td>
+        <td></td>
+    </tr>
+    <tr>
+        <td>paddle版本</td>
+        <td></td>
+        <td></td>
+    </tr>
+</table>
+上表中的配置是基于4卡GPU训练的。
+**接下来** 以 `FixMatch_CCSSL_cifar10_40000_4gpu.yaml` 配置和训练好的模型文件为例，展示在cifar10数据集上进行训练，测试，推理的过程。
+
+## 3. 数据准备
+cifar10数据在训练过程中会自动下载到默认缓存路径 `~/.cache/paddle/dataset/cifar/`中，请在训练过程中保持网络畅通。
+
+## 4. 模型训练
+1. 执行以下命令开始训练
+
+单卡训练执行以下命令
+```
+python tools/train.py -c ppcls/configs/ssl/FixMatch_CCSSL_cifar10_4000.yaml
+```
+
+4卡训练执行以下操作
+```
+python -m paddle.distributed.launch --gpus='0,1,2,3' tools/train.py -c ppcls/configs/ssl/FixMatchCCSSL/FixMatchCCSSL_cifar10_4000_4gpu.yaml
+```
+
+2. **查看训练日志和保存的模型参数文件** 训练过程中屏幕会实时打印loss等指标信息，同时会保存日志文件 `train.log` ，模型参数文件 `*.pdparams`，优化器参数文件 `*.pdopt` 等内容到`Global.output_dir`指定的文件夹下，默认在 `PaddleClas/output/WideResNet/`文件夹下。
+
+## 5. 模型评估与推理部署
+### 5.1 模型评估
+准备用于评估的 `*.pdparams` 模型参数文件，可以使用训练好的模型，也可以使用 *4. 模型训练* 中保存的模型。
+* 以训练过程中保存的 `best_model_ema.ema.pdparams`为例，执行如下命令即可进行评估。
+```
+python3.7 tools/eval.py -c ppcls/configs/ssl/FixMatch_CCSSL_cifar10_4000.yaml -o Global.pretrained_model="./output/WideResNetCCSSL/best_model_ema.ema"
+```
+
+* 以训练好的模型为例，下载提供的已经训练好的模型，到 `PaddleClas/pretrained_models` 文件夹中，执行如下命令即可进行评估。
+```
+# 下载模型
+cd PaddleClas
+mkdir pretrained_models
+cd pretrained_models
+wget 
+cd ..
+# 评估
+python3.7 tools/eval.py -c ppcls/configs/ssl/FixMatch_CCSSL_cifar10_4000.yaml -o Global.pretrained_model=""
+```
+**注：** `pretrained_model` 后填入的地址不需要加 `.pdparams`后缀，在程序运行时会自动补上。
+
+* 查看输出结果
+```
+[2022/12/08 09:36:13] ppcls INFO: [Eval][Epoch 0][Iter: 0/157]CELoss: 0.00999, loss: 0.00999, top1: 1.00000, top5: 1.00000, batch_cost: 5.11046s, reader_cost: 1.22196, ips: 12.52334 images/sec
+[2022/12/08 09:36:13] ppcls INFO: [Eval][Epoch 0][Iter: 20/157]CELoss: 0.04825, loss: 0.04825, top1: 0.95164, top5: 1.00000, batch_cost: 0.02071s, reader_cost: 0.00207, ips: 3089.66447 images/sec
+[2022/12/08 09:36:14] ppcls INFO: [Eval][Epoch 0][Iter: 40/157]CELoss: 0.03500, loss: 0.03500, top1: 0.95084, top5: 1.00000, batch_cost: 0.02155s, reader_cost: 0.00108, ips: 2970.07129 images/sec
+[2022/12/08 09:36:14] ppcls INFO: [Eval][Epoch 0][Iter: 60/157]CELoss: 0.26421, loss: 0.26421, top1: 0.94928, top5: 0.99949, batch_cost: 0.02048s, reader_cost: 0.00151, ips: 3124.81965 images/sec
+[2022/12/08 09:36:14] ppcls INFO: [Eval][Epoch 0][Iter: 80/157]CELoss: 0.16254, loss: 0.16254, top1: 0.95332, top5: 0.99942, batch_cost: 0.02124s, reader_cost: 0.00117, ips: 3013.43961 images/sec
+[2022/12/08 09:36:15] ppcls INFO: [Eval][Epoch 0][Iter: 100/157]CELoss: 0.15471, loss: 0.15471, top1: 0.95374, top5: 0.99923, batch_cost: 0.02046s, reader_cost: 0.00098, ips: 3128.15428 images/sec
+[2022/12/08 09:36:15] ppcls INFO: [Eval][Epoch 0][Iter: 120/157]CELoss: 0.05237, loss: 0.05237, top1: 0.95493, top5: 0.99935, batch_cost: 0.02061s, reader_cost: 0.00084, ips: 3106.03867 images/sec
+[2022/12/08 09:36:16] ppcls INFO: [Eval][Epoch 0][Iter: 140/157]CELoss: 0.03242, loss: 0.03242, top1: 0.95601, top5: 0.99945, batch_cost: 0.02084s, reader_cost: 0.00075, ips: 3071.00311 images/sec
+[2022/12/08 09:36:16] ppcls INFO: [Eval][Epoch 0][Avg]CELoss: 0.16041, loss: 0.16041, top1: 0.95610, top5: 0.99950
+```
+默认评估日志保存在 `PaddleClas/output/WideResNetCCSSL/eval.log`中，可以看到我们提供的模型在cifar10数据集上的评估指标为top1: 95.61, top5: 99.95
+
+### 5.2 模型推理
+#### 5.2.1 推理模型准备
+将训练过程中保存的模型文件转成inference模型，同样以 `best_model_ema.ema_pdparams`为例，执行以下命令进行转换
+```
+python3.7 tools/export_model.py \
+-c ppcls/configs/ssl/FixMatch_CCSSL/FixMatch_CCSSL_cifar10_4000.yaml \
+-o Global.pretrained_model="output/WideResNetCCSSL/best_model_ema.ema" \
+-o Global.save_inference_fir="./deploy/inference"
+```
+
+#### 5.2.2 基于 Python 预测引擎推理
+1. 修改 `PaddleClas/deploy/configs/inference_cls.yaml`
+* * 将`infer_imgs:` 后的路径段改为 query 文件夹下的任意一张图片路径（下方配置使用的是`demo.jpg`图片的路径）
+* * 将`rec_inference_model.dir:` 后的字段改为解压出来的 inference 模型文件夹路径
+* * 将`transform_ops:` 字段下的预处理配置改为 `FixMatch_CCSSL_cifar10_40000.yaml` 中 `Eval.dataset`下的预处理配置
+```
+Global:
+  infer_imgs: "demo"
+  rec_inference_model_dir: "./inferece"
+  batch_size: 1
+  use_gpu: False
+  enable_mkldnn: True
+  cpu_num_threads: 10
+  enable_benchmark: False
+  use_fp16: False
+  ir_optim: True
+  use_tensorrt: False
+  gpu_mem: 8000
+  enable_profile: False
+
+RecPreProcess:
+  transform_ops:
+   -  NormalizeImage:
+        scale: 1.0/255.0
+        mean: [0.4914, 0.4822, 0.4465]
+        std: [0.2471, 0.2435, 0.2616]
+        order: hwc
+PostProcess: null
+```
+
+2. 执行推理命令
+```
+cd ./deploy/
+python3.7 python/predict_rec.py -c ./configs/inference_rec.yaml
+```
+
+3. 查看输出结果，实际结果为一个长度为10的向量，表示图像分类的结果，如
+```
+
+```
+
+#### 5.2.3 基于C++预测引擎推理
+PaddleClas 提供了基于 C++ 预测引擎推理的示例，您可以参考[服务器端 C++ 预测](https://github.com/zh-hike/PaddleClas/blob/develop/docs/zh_CN/deployment/image_classification/cpp/linux.md)来完成相应的推理部署。如果您使用的是 Windows 平台，可以参考基于 Visual Studio 2019 Community CMake 编译指南完成相应的预测库编译和模型预测工作。
+
+### 5.4 服务化部署
+Paddle Serving 提供高性能、灵活易用的工业级在线推理服务。Paddle Serving 支持 RESTful、gRPC、bRPC 等多种协议，提供多种异构硬件和多种操作系统环境下推理解决方案。更多关于Paddle Serving 的介绍，可以参考Paddle Serving 代码仓库。
+
+## 5.5 端侧部署
+Paddle Lite 是一个高性能、轻量级、灵活性强且易于扩展的深度学习推理框架，定位于支持包括移动端、嵌入式以及服务器端在内的多硬件平台。更多关于 Paddle Lite 的介绍，可以参考Paddle Lite 代码仓库。
+
+PaddleClas 提供了基于 Paddle Lite 来完成模型[端侧部署](https://github.com/zh-hike/PaddleClas/blob/develop/docs/zh_CN/deployment/image_classification/paddle_lite.md)的示例，您可以参考端侧部署来完成相应的部署工作。
+
+## Paddle2ONNX 模型转换与预测
+Paddle2ONNX 支持将 PaddlePaddle 模型格式转化到 ONNX 模型格式。通过 ONNX 可以完成将 Paddle 模型到多种推理引擎的部署，包括TensorRT/OpenVINO/MNN/TNN/NCNN，以及其它对 ONNX 开源格式进行支持的推理引擎或硬件。更多关于 Paddle2ONNX 的介绍，可以参考Paddle2ONNX 代码仓库。
+
+PaddleClas 提供了基于 Paddle2ONNX 来完成 inference 模型转换 ONNX 模型并作推理预测的示例，您可以参考 [Paddle2ONNX](https://github.com/zh-hike/PaddleClas/blob/develop/docs/zh_CN/deployment/image_classification/paddle2onnx.md) 模型转换与预测来完成相应的部署工作。
+
+## 6. 参考资料
+1. [CCSSL](https://arxiv.org/abs/2203.02261)
\ No newline at end of file

ppcls/arch/backbone/model_zoo/wideresnet.py

@@ -201,7 +201,9 @@ class Wide_ResNet(nn.Layer):
         feat = self.relu(self.bn1(feat))
         feat = F.adaptive_avg_pool2d(feat, 1)
         feat = paddle.reshape(feat, [-1, self.channels])
-
+        if not self.training:
+            return self.fc(feat)
+            
         if self.proj:
             pfeat = self.fc1(feat)
             pfeat = self.relu_mlp(pfeat)

ppcls/configs/ssl/FixMatchCCSSL/FixMatchCCSSL_cifar10_4000_4gpu.yaml

+Global:
+  checkpoints: null
+  pretrained_model: null
+  output_dir: ./output
+  device: gpu
+  save_interval: -1
+  eval_during_train: true
+  eval_interval: 1
+  epochs: 1024
+  iter_per_epoch: 200
+  print_batch_step: 20
+  use_visualdl: false
+  use_dali: false
+  train_mode: fixmatch_ccssl
+  image_shape: [3, 32, 32]
+  save_inference_dir: ./inference
+
+SSL:
+  T: 1
+  threshold: 0.95
+
+EMA:
+  decay: 0.999
+
+Arch:
+  name: WideResNet
+  widen_factor: 2
+  depth: 28
+  dropout: 0     # CCSSL为 drop_rate
+  num_classes: &sign_num_classes 10
+  low_dim: 64
+  proj: true
+  proj_after: false
+
+  use_sync_bn: true
+
+Loss:
+  Train:
+    - CELoss:
+        weight: 1.0
+        reduction: "mean"
+  Eval:
+    - CELoss:
+        weight: 1.0
+  
+UnLabelLoss:
+  Train:
+    - CCSSLCeLoss:
+        weight: 1.
+    - SoftSupConLoss:
+        weight: 1.0
+        temperature: 0.07
+    # - CCSSLLoss:
+    #     CELoss:
+    #       weight: 1.0
+    #       reduction: "none"
+    #     SoftSupConLoss:
+    #       weight: 1.0
+    #       temperature: 0.07
+    #     weight: 1.
+
+Optimizer:
+  name: Momentum
+  momentum: 0.9
+  use_nesterov: true
+  weight_decay: 0.001
+  lr:
+    name: 'cosine_schedule_with_warmup'
+    learning_rate: 0.03
+    num_warmup_steps: 0
+    num_training_steps: 524800
+
+DataLoader:
+  mean: &sign_mean [0.4914, 0.4822, 0.4465]
+  std: &sign_std [0.2471, 0.2435, 0.2616]
+  Train:
+    dataset:
+      name: CIFAR10SSL
+      data_file: null
+      mode: 'train'
+      download: true
+      sample_per_label: 400
+      expand_labels: 1
+      transform_ops:
+        - RandomHorizontalFlip:
+            prob: 0.5
+        - RandomCrop:
+            size: 32
+            padding: 4
+            padding_mode: "reflect"
+        - ToTensor:
+        - Normalize:
+            mean: *sign_mean
+            std: *sign_std
+    
+    sampler:
+      name: DistributedBatchSampler    # DistributedBatchSampler
+      batch_size: 64
+      drop_last: true
+      shuffle: true
+    loader:
+      num_workers: 4
+      use_shared_memory: true
+  
+
+
+  UnLabelTrain:
+    dataset:
+      name: CIFAR10SSL
+      data_file: null
+      mode: 'train'
+      download: true
+
+      transform_w:
+        - RandomHorizontalFlip:
+            prob: 0.5
+        - RandomCrop:
+            size: 32
+            padding: 4
+            padding_mode: 'reflect'
+        - ToTensor:
+        - Normalize:
+            mean: *sign_mean
+            std: *sign_std
+
+      transform_s1:
+        - RandomHorizontalFlip:
+                prob: 0.5
+        - RandomCrop:
+            size: 32
+            padding: 4
+            padding_mode: 'reflect'
+        - RandAugmentMC:
+            n: 2
+            m: 10
+        - ToTensor:
+        - Normalize:
+            mean: *sign_mean
+            std: *sign_std
+      
+      transform_s2:
+        - RandomResizedCrop:
+            size: 32
+        - RandomHorizontalFlip:
+            prob: 0.5
+        - RandomApply:
+            transforms:
+              - ColorJitter:
+                  brightness: 0.4
+                  contrast: 0.4
+                  saturation: 0.4
+                  hue: 0.1
+            p: 0.8
+        - RandomGrayscale:
+            p: 0.2
+        - ToTensor:
+        # - Normalize:
+        #     mean: *sign_mean
+        #     std: *sign_std
+
+    sampler:
+      name: DistributedBatchSampler   # DistributedBatchSampler
+      batch_size: 448
+      drop_last: true
+      shuffle: true
+    loader:
+      num_workers: 4
+      use_shared_memory: true
+
+  Eval:
+    dataset:
+      name: CIFAR10SSL
+      mode: 'test'
+      download: true
+      data_file: null
+      transform_ops:
+        - ToTensor:
+        - Normalize:
+            mean: *sign_mean
+            std: *sign_std
+    sampler:
+      name: DistributedBatchSampler
+      batch_size: 64
+      drop_last: False
+      shuffle: True
+    loader:
+      num_workers: 4
+      use_shared_memory: true
+
+
+
+Metric:
+  Eval:
+    - TopkAcc:
+        topk: [1, 5]
\ No newline at end of file

ppcls/data/__init__.py

@@ -35,7 +35,7 @@ from ppcls.data.dataloader.multi_scale_dataset import MultiScaleDataset
 from ppcls.data.dataloader.person_dataset import Market1501, MSMT17
 from ppcls.data.dataloader.face_dataset import FiveValidationDataset, AdaFaceDataset
 from ppcls.data.dataloader.custom_label_dataset import CustomLabelDataset
-from ppcls.data.dataloader.cifar import Cifar10, Cifar100
+from ppcls.data.dataloader.cifar import Cifar10, Cifar100, CIFAR10SSL
 
 # sampler
 from ppcls.data.dataloader.DistributedRandomIdentitySampler import DistributedRandomIdentitySampler

ppcls/data/dataloader/__init__.py

@@ -13,3 +13,4 @@ from ppcls.data.dataloader.person_dataset import Market1501, MSMT17
 from ppcls.data.dataloader.face_dataset import AdaFaceDataset, FiveValidationDataset
 from ppcls.data.dataloader.custom_label_dataset import CustomLabelDataset
 from ppcls.data.dataloader.cifar import Cifar10, Cifar100
+from ppcls.data.dataloader.cifar import CIFAR10SSL

ppcls/data/dataloader/cifar.py

@@ -15,11 +15,16 @@
 from __future__ import print_function
 import numpy as np
 import cv2
+import shutil
 from ppcls.data import preprocess
 from ppcls.data.preprocess import transform
+from ppcls.data.preprocess import BaseTransform, ListTransform
 from ppcls.data.dataloader.common_dataset import create_operators
 from paddle.vision.datasets import Cifar10 as Cifar10_paddle
 from paddle.vision.datasets import Cifar100 as Cifar100_paddle
+from paddle.vision.datasets import cifar
+import os
+from PIL import Image
 
 
 class Cifar10(Cifar10_paddle):
@@ -112,4 +117,105 @@ class Cifar100(Cifar100_paddle):
             image3 = transform(image, self._transform_ops_strong)
             image3 = image3.transpose((2, 0, 1))
 
-            return (image2, image3, np.int64(label))
\ No newline at end of file
+            return (image2, image3, np.int64(label))
+
+
+def np_convert_pil(array):
+    """
+    array conver image
+    Args:
+        array: array and dim is 1
+    """
+    assert len(array.shape), "dim of array should 1"
+    img = Image.fromarray(array.reshape(3, 32, 32).transpose(1, 2, 0))
+    return img
+
+
+class CIFAR10(cifar.Cifar10):
+    """
+    cifar10 dataset
+    """
+    def __init__(self, data_file, download=True, mode='train'):
+        super().__init__(download=download, mode=mode)
+        if data_file is not None:
+            os.makedirs(data_file, exist_ok=True)
+            if not os.path.exists(os.path.join(data_file, 'cifar-10-python.tar.gz')):
+                shutil.move('~/.cache/paddle/dataset/cifar/cifar-10-python.tar.gz', data_file)
+        self.num_classes = 10
+        self.x = []
+        self.y = []
+        for d in self.data:
+            self.x.append(d[0])
+            self.y.append(d[1])
+
+        self.x = np.array(self.x)
+        self.y = np.array(self.y)
+
+    def __getitem__(self, idx):
+        return self.x[idx], self.y[idx]
+
+    def __len__(self):
+        return self.x.shape[0]
+
+
+class CIFAR10SSL(CIFAR10):
+    """
+    from Cifar10
+    """
+
+    def __init__(self, 
+                 data_file=None, 
+                 sample_per_label=None, 
+                 download=True,
+                 expand_labels=1,
+                 mode='train',
+                 transform_ops=None,
+                 transform_w=None,
+                 transform_s1=None,
+                 transform_s2=None):
+        super().__init__(data_file, download=download, mode=mode)
+        self.data_type = 'unlabeled_train' if mode == 'train' else 'val'
+        if transform_ops is not None and sample_per_label is not None:
+            index = []
+            self.data_type = 'labeled_train'
+            for c in range(self.num_classes):
+                idx = np.where(self.y == c)[0]
+                idx = np.random.choice(idx, sample_per_label, False)
+                index.extend(idx)
+            index = index * expand_labels
+            # print(index)
+            self.x = self.x[index]
+            self.y = self.y[index]
+        self.transforms = [transform_ops] if transform_ops is not None else [transform_w, transform_s1, transform_s2]
+        self.mode = mode
+
+    def __getitem__(self, idx):
+        img, label = np_convert_pil(self.x[idx]), self.y[idx]
+        results = ListTransform(self.transforms)(img)
+        if self.data_type == 'unlabeled_train':
+            return results
+        return results[0], label
+        
+    def __len__(self):
+        return self.x.shape[0]
+        
+
+# def x_u_split(num_labeled, num_classes, label):
+#     """
+#     split index of dataset to labeled x and unlabeled u
+#     Args:
+#         num_labeled: num of labeled dataset
+#         label: list or array, label
+#     """
+#     assert num_labeled <= len(label), "arg num_labeled should <= num of label"
+#     label = np.array(label) if isinstance(label, list) else label
+#     label_per_class = num_labeled // num_classes
+#     labeled_idx = []
+#     unlabeled_idx = np.array(list(range(label.shape[0])))
+#     for c in range(num_classes):
+#         idx = np.where(label == c)[0]
+#         idx = np.random.choice(idx, label_per_class, False)
+#         labeled_idx.extend(idx)
+    
+#     np.random.shuffle(labeled_idx)
+#     return labeled_idx, unlabeled_idx
\ No newline at end of file

ppcls/data/preprocess/__init__.py

@@ -39,6 +39,7 @@ from ppcls.data.preprocess.ops.operators import RandomResizedCrop
 from ppcls.data.preprocess.ops.operators import CropWithPadding
 from ppcls.data.preprocess.ops.operators import RandomInterpolationAugment
 from ppcls.data.preprocess.ops.operators import ColorJitter
+from ppcls.data.preprocess.ops.operators import RandomGrayscale
 from ppcls.data.preprocess.ops.operators import RandomCropImage
 from ppcls.data.preprocess.ops.operators import RandomRotation
 from ppcls.data.preprocess.ops.operators import Padv2
@@ -49,9 +50,13 @@ from paddle.vision.transforms import Pad as Pad_paddle_vision
 from ppcls.data.preprocess.batch_ops.batch_operators import MixupOperator, CutmixOperator, OpSampler, FmixOperator
 from ppcls.data.preprocess.batch_ops.batch_operators import MixupCutmixHybrid
 
+from .ops.randaugmentmc import RandAugmentMC, RandomApply
+
 import numpy as np
 from PIL import Image
 import random
+from paddle.vision.transforms import transforms as T
+from paddle.vision.transforms.transforms import RandomCrop, ToTensor, Normalize
 
 
 def transform(data, ops=[]):
@@ -117,3 +122,41 @@ class TimmAutoAugment(RawTimmAutoAugment):
             img = np.asarray(img)
 
         return img
+
+
+class BaseTransform:
+    def __init__(self, cfg) -> None:
+        """
+        Args:
+            cfg: list [dict, dict, dict]
+        """
+        ts = []
+        for op in cfg:
+            name = list(op.keys())[0]
+            if op[name] is None:
+                ts.append(eval(name)())
+            else:
+                ts.append(eval(name)(**(op[name])))
+
+        self.t = T.Compose(ts)
+
+    def __call__(self, img):
+        
+        return self.t(img)
+
+
+class ListTransform:
+    def __init__(self, ops) -> None:
+        """
+        Args:
+            ops: list[list[dict, dict], ...]
+        """
+        self.ts = []
+        for op in ops:
+            self.ts.append(BaseTransform(op))
+
+    def __call__(self, img):
+        results = []
+        for op in self.ts:
+            results.append(op(img))
+        return results

ppcls/data/preprocess/ops/operators.py

@@ -835,3 +835,37 @@ class BlurImage(object):
                                         self.motion_max_angle)
             label = 1
         return {"img": img, "blur_image": label}
+
+
+class RandomGrayscale(object):
+    """Randomly convert image to grayscale with a probability of p (default 0.1).
+
+    Args:
+        p (float): probability that image should be converted to grayscale.
+
+    Returns:
+        PIL Image: Grayscale version of the input image with probability p and unchanged
+        with probability (1-p).
+        - If input image is 1 channel: grayscale version is 1 channel
+        - If input image is 3 channel: grayscale version is 3 channel with r == g == b
+
+    """
+
+    def __init__(self, p=0.1):
+        self.p = p
+
+    def __call__(self, img):
+        """
+        Args:
+            img (PIL Image): Image to be converted to grayscale.
+
+        Returns:
+            PIL Image: Randomly grayscaled image.
+        """
+        num_output_channels = 1 if img.mode == 'L' else 3
+        if random.random() < self.p:
+            return F.to_grayscale(img, num_output_channels=num_output_channels)
+        return img
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
\ No newline at end of file

ppcls/data/preprocess/ops/randaugmentmc.py

+import PIL
+import PIL.ImageDraw
+import random
+from paddle.vision.transforms import transforms as T
+from paddle.vision.transforms.transforms import ColorJitter
+import numpy as np
+
+
+PARAMETER_MAX = 10
+
+
+
+def AutoContrast(img, **kwarg):
+    return PIL.ImageOps.autocontrast(img)
+
+
+def Brightness(img, v, max_v, bias=0):
+    v = _float_parameter(v, max_v) + bias
+    return PIL.ImageEnhance.Brightness(img).enhance(v)
+
+
+def Color(img, v, max_v, bias=0):
+    v = _float_parameter(v, max_v) + bias
+    return PIL.ImageEnhance.Color(img).enhance(v)
+
+def Contrast(img, v, max_v, bias=0):
+    v = _float_parameter(v, max_v) + bias
+    return PIL.ImageEnhance.Contrast(img).enhance(v)
+
+
+def Cutout(img, v, max_v, bias=0):
+    if v == 0:
+        return img
+    v = _float_parameter(v, max_v) + bias
+    v = int(v * min(img.size))
+    return CutoutAbs(img, v)
+
+
+def CutoutAbs(img, v, **kwarg):
+    w, h = img.size
+    x0 = np.random.uniform(0, w)
+    y0 = np.random.uniform(0, h)
+    x0 = int(max(0, x0 - v / 2.))
+    y0 = int(max(0, y0 - v / 2.))
+    x1 = int(min(w, x0 + v))
+    y1 = int(min(h, y0 + v))
+    xy = (x0, y0, x1, y1)
+    # gray
+    color = (127, 127, 127)
+    img = img.copy()
+    PIL.ImageDraw.Draw(img).rectangle(xy, color)
+    return img
+
+
+def Equalize(img, **kwarg):
+    return PIL.ImageOps.equalize(img)
+
+
+def Identity(img, **kwarg):
+    return img
+
+
+def Invert(img, **kwarg):
+    return PIL.ImageOps.invert(img)
+
+
+def Posterize(img, v, max_v, bias=0):
+    v = _int_parameter(v, max_v) + bias
+    return PIL.ImageOps.posterize(img, v)
+
+
+def Rotate(img, v, max_v, bias=0):
+    v = _int_parameter(v, max_v) + bias
+    if random.random() < 0.5:
+        v = -v
+    return img.rotate(v)
+
+
+def Sharpness(img, v, max_v, bias=0):
+    v = _float_parameter(v, max_v) + bias
+    return PIL.ImageEnhance.Sharpness(img).enhance(v)
+
+
+def ShearX(img, v, max_v, bias=0):
+    v = _float_parameter(v, max_v) + bias
+    if random.random() < 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.Transform.AFFINE, (1, v, 0, 0, 1, 0))
+
+
+def ShearY(img, v, max_v, bias=0):
+    v = _float_parameter(v, max_v) + bias
+    if random.random() < 0.5:
+        v = -v
+    return img.transform(img.size, PIL.Image.Transform.AFFINE, (1, 0, 0, v, 1, 0))
+
+
+def Solarize(img, v, max_v, bias=0):
+    v = _int_parameter(v, max_v) + bias
+    return PIL.ImageOps.solarize(img, 256 - v)
+
+
+def SolarizeAdd(img, v, max_v, bias=0, threshold=128):
+    v = _int_parameter(v, max_v) + bias
+    if random.random() < 0.5:
+        v = -v
+    img_np = np.array(img).astype(np.int)
+    img_np = img_np + v
+    img_np = np.clip(img_np, 0, 255)
+    img_np = img_np.astype(np.uint8)
+    img = Image.fromarray(img_np)
+    return PIL.ImageOps.solarize(img, threshold)
+
+
+def TranslateX(img, v, max_v, bias=0):
+    v = _float_parameter(v, max_v) + bias
+    if random.random() < 0.5:
+        v = -v
+    v = int(v * img.size[0])
+    return img.transform(img.size, PIL.Image.Transform.AFFINE, (1, 0, v, 0, 1, 0))
+
+
+def TranslateY(img, v, max_v, bias=0):
+    v = _float_parameter(v, max_v) + bias
+    if random.random() < 0.5:
+        v = -v
+    v = int(v * img.size[1])
+    return img.transform(img.size, PIL.Image.Transform.AFFINE, (1, 0, 0, 0, 1, v))
+
+
+def _float_parameter(v, max_v):
+    return float(v) * max_v / PARAMETER_MAX
+
+
+def _int_parameter(v, max_v):
+    return int(v * max_v / PARAMETER_MAX)
+
+
+def fixmatch_augment_pool():
+    # FixMatch paper
+    augs = [(AutoContrast, None, None),
+            (Brightness, 0.9, 0.05),
+            (Color, 0.9, 0.05),
+            (Contrast, 0.9, 0.05),
+            (Equalize, None, None),
+            (Identity, None, None),
+            (Posterize, 4, 4),
+            (Rotate, 30, 0),
+            (Sharpness, 0.9, 0.05),
+            (ShearX, 0.3, 0),
+            (ShearY, 0.3, 0),
+            (Solarize, 256, 0),
+            (TranslateX, 0.3, 0),
+            (TranslateY, 0.3, 0)]
+    return augs
+
+
+def my_augment_pool():
+    # Test
+    augs = [(AutoContrast, None, None),
+            (Brightness, 1.8, 0.1),
+            (Color, 1.8, 0.1),
+            (Contrast, 1.8, 0.1),
+            (Cutout, 0.2, 0),
+            (Equalize, None, None),
+            (Invert, None, None),
+            (Posterize, 4, 4),
+            (Rotate, 30, 0),
+            (Sharpness, 1.8, 0.1),
+            (ShearX, 0.3, 0),
+            (ShearY, 0.3, 0),
+            (Solarize, 256, 0),
+            (SolarizeAdd, 110, 0),
+            (TranslateX, 0.45, 0),
+            (TranslateY, 0.45, 0)]
+    return augs
+
+
+class RandAugmentPC(object):
+    def __init__(self, n, m):
+        assert n >= 1
+        assert 1 <= m <= 10
+        self.n = n
+        self.m = m
+        self.augment_pool = my_augment_pool()
+
+    def __call__(self, img):
+        ops = random.choices(self.augment_pool, k=self.n)
+        for op, max_v, bias in ops:
+            prob = np.random.uniform(0.2, 0.8)
+            if random.random() + prob >= 1:
+                img = op(img, v=self.m, max_v=max_v, bias=bias)
+        img = CutoutAbs(img, int(32*0.5))
+        return img
+
+
+class RandAugmentMC(object):
+    def __init__(self, n, m):
+        assert n >= 1
+        assert 1 <= m <= 10
+        self.n = n
+        self.m = m
+        self.augment_pool = fixmatch_augment_pool()
+
+    def __call__(self, img):
+        ops = random.choices(self.augment_pool, k=self.n)
+        for op, max_v, bias in ops:
+            v = np.random.randint(1, self.m)
+            if random.random() < 0.5:
+                img = op(img, v=v, max_v=max_v, bias=bias)
+        img = CutoutAbs(img, int(32 * 0.5))
+        return img
+
+
+class RandomApply:
+    def __init__(self, p, transforms):
+        self.p = p
+        ts = []
+        for t in transforms:
+            for key in t.keys():
+                ts.append(eval(key)(**t[key]))
+
+        self.trans = T.Compose(ts)
+
+    def __call__(self, img):
+        timg = self.trans(img)
+        return timg
\ No newline at end of file

ppcls/engine/train/__init__.py

@@ -12,4 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from ppcls.engine.train.train import train_epoch
-from ppcls.engine.train.train_fixmatch import train_epoch_fixmatch
\ No newline at end of file
+from ppcls.engine.train.train_fixmatch import train_epoch_fixmatch
+from ppcls.engine.train.train_fixmatch_ccssl import train_epoch_fixmatch_ccssl
\ No newline at end of file

ppcls/engine/train/train_fixmatch_ccssl.py

+
+
+from __future__ import absolute_import, division, print_function
+
+import time
+from turtle import update
+import paddle
+from ppcls.engine.train.train_fixmatch import get_loss
+from ppcls.engine.train.utils import update_loss, update_metric, log_info
+from ppcls.utils import profiler
+from paddle.nn import functional as F
+import numpy as np
+# from reprod_log import ReprodLogger
+
+
+def train_epoch_fixmatch_ccssl(engine, epoch_id, print_batch_step):
+    ##############################################################
+    # out_logger = ReprodLogger()
+    # loss_logger = ReprodLogger()
+    # epoch = 0
+    ##############################################################
+
+    tic = time.time()
+    if not hasattr(engine, 'train_dataloader_iter'):
+        engine.train_dataloader_iter = iter(engine.train_dataloader)
+        engine.unlabel_train_dataloader_iter = iter(engine.unlabel_train_dataloader)
+    
+    temperture = engine.config['SSL'].get("T", 1)
+    threshold = engine.config['SSL'].get("threshold", 0.95)
+    assert engine.iter_per_epoch is not None, "Global.iter_per_epoch need to be set"
+    threshold = paddle.to_tensor(threshold)
+
+    # dataload_logger = ReprodLogger()
+
+    for iter_id in range(engine.iter_per_epoch):
+        if iter_id >= engine.iter_per_epoch:
+            break
+
+        if iter_id == 5:
+            for key in engine.time_info:
+                engine.time_info[key].reset()
+
+        try:
+            label_data_batch = engine.train_dataloader_iter.next()
+        except Exception:
+            engine.train_dataloader_iter = iter(engine.train_dataloader)
+            label_data_batch = engine.train_dataloader_iter.next()
+
+        try:
+            unlabel_data_batch = engine.unlabel_train_dataloader_iter.next()
+        except Exception:
+            engine.unlabel_train_dataloader_iter = iter(engine.unlabel_train_dataloader)
+            unlabel_data_batch = engine.unlabel_train_dataloader_iter.next()
+
+        assert len(unlabel_data_batch) == 3
+        assert unlabel_data_batch[0].shape == unlabel_data_batch[1].shape == unlabel_data_batch[2].shape
+
+        ##############################################################
+        # inputs_x, target_x = label_data_batch
+        # inputs_w, inputs_s1, inputs_s2 = unlabel_data_batch
+        # dataload_logger.add(f"inputs_x_iter_{iter_id}", inputs_x.detach().numpy())
+        # dataload_logger.add(f"target_x_iter_{iter_id}", target_x.detach().numpy())
+        # dataload_logger.add(f"inputs_w_iter_{iter_id}", inputs_w.detach().numpy())
+        # dataload_logger.add(f"inputs_s1_iter_{iter_id}", inputs_s1.detach().numpy())
+        # dataload_logger.add(f"inputs_s2_iter_{iter_id}", inputs_s2.detach().numpy())
+        # dataload_logger.save('../align/step2/data/paddle.npy')
+
+        # assert 1==0
+
+        ##############################################################
+        engine.time_info['reader_cost'].update(time.time() - tic)
+        batch_size = label_data_batch[0].shape[0] \
+                    + unlabel_data_batch[0].shape[0] \
+                    + unlabel_data_batch[1].shape[0] \
+                    + unlabel_data_batch[2].shape[0]
+        engine.global_step += 1
+
+        # make inputs
+        inputs_x, targets_x = label_data_batch
+        # inputs_x = inputs_x[0]
+        inputs_w, inputs_s1, inputs_s2 = unlabel_data_batch
+
+        batch_size_label = inputs_x.shape[0]
+        inputs = paddle.concat([inputs_x, inputs_w, inputs_s1, inputs_s2])
+
+        loss_dict, logits_label = get_loss(engine, inputs, batch_size_label, 
+                                           temperture, threshold, targets_x,
+                                        #    epoch=epoch,
+                                        #    batch_idx=iter_id,
+                                        #    out_logger=out_logger,
+                                        #    loss_logger=loss_logger
+                                           )
+
+        loss = loss_dict['loss']
+        loss.backward()
+
+        
+        for i in range(len(engine.optimizer)):
+            engine.optimizer[i].step()
+        
+        for i in range(len(engine.lr_sch)):
+            if not getattr(engine.lr_sch[i], 'by_epoch', False):
+                engine.lr_sch[i].step()
+
+        for i in range(len(engine.optimizer)):
+            engine.optimizer[i].clear_grad()
+
+        if engine.ema:
+            engine.model_ema.update(engine.model)
+        
+        update_metric(engine, logits_label, label_data_batch, batch_size)
+
+        update_loss(engine, loss_dict, batch_size)
+
+        engine.time_info['batch_cost'].update(time.time() - tic)
+
+        if iter_id % print_batch_step == 0:
+            log_info(engine, batch_size, epoch_id, iter_id)
+
+        tic = time.time()
+
+        # if iter_id == 10:
+        #     assert 1==0
+
+    for i in range(len(engine.lr_sch)):
+        if getattr(engine.lr_sch[i], 'by_epoch', False):
+            engine.lr_sch[i].step()
+
+def get_loss(engine,
+             inputs,
+             batch_size_label,
+             temperture,
+             threshold,
+             targets_x,
+             **kwargs
+             ):
+
+    logits, feats = engine.model(inputs)
+    feat_w, feat_s1, feat_s2 = feats[batch_size_label:].chunk(3)
+    feat_x = feats[:batch_size_label]
+    logits_x = logits[:batch_size_label]
+    logits_w, logits_s1, logits_s2 = logits[batch_size_label:].chunk(3)
+    loss_dict_label = engine.train_loss_func(logits_x, targets_x)
+    probs_u_w = F.softmax(logits_w.detach() / temperture, axis=-1)
+    max_probs, p_targets_u_w = probs_u_w.max(axis=-1), probs_u_w.argmax(axis=-1)
+    mask = paddle.greater_equal(max_probs, threshold).astype('float')
+    
+    # feats = paddle.concat([logits_s1.unsqueeze(1), logits_s2.unsqueeze(1)], axis=1)
+    feats = paddle.concat([feat_s1.unsqueeze(1), feat_s2.unsqueeze(1)], axis=1)
+    batch = {'logits_w': logits_w,
+             'logits_s1': logits_s1,
+             'p_targets_u_w': p_targets_u_w,
+             'mask': mask,
+             'max_probs': max_probs,
+             }
+
+    unlabel_loss = engine.unlabel_train_loss_func(feats, batch)
+
+    loss_dict = {}
+    for k, v in loss_dict_label.items():
+        if k != 'loss':
+            loss_dict[k] = v
+    for k, v in unlabel_loss.items():
+        if k != 'loss':
+            loss_dict[k] = v
+    loss_dict['loss'] = loss_dict_label['loss'] + unlabel_loss['loss']
+
+    ##############################################################
+    # print(loss_dict)
+    # epoch = kwargs['epoch']
+    # batch_idx = kwargs['batch_idx']
+    # out_logger = kwargs['out_logger']
+    # loss_logger = kwargs['loss_logger']
+    # out_logger.add(f'logit_x_{epoch}_{batch_idx}', logits_x.detach().numpy())
+    # out_logger.add(f'logit_u_w_{epoch}_{batch_idx}', logits_w.detach().numpy())
+    # out_logger.add(f'logit_u_s1_{epoch}_{batch_idx}', logits_s1.detach().numpy())
+    # out_logger.add(f'logit_u_s2_{epoch}_{batch_idx}', logits_s2.detach().numpy())
+
+    # out_logger.add(f'feat_x_{epoch}_{batch_idx}', feat_x.detach().numpy())
+    # out_logger.add(f'feat_w_{epoch}_{batch_idx}', feat_w.detach().numpy())
+    # out_logger.add(f'feat_s1_{epoch}_{batch_idx}', feat_s1.detach().numpy())
+    # out_logger.add(f'feat_s2_{epoch}_{batch_idx}', feat_s2.detach().numpy())
+
+    # loss_logger.add(f'loss_{epoch}_{batch_idx}', loss_dict['loss'].detach().numpy())
+    # loss_logger.add(f'loss_x_{epoch}_{batch_idx}', loss_dict['CELoss'].detach().cpu().numpy())
+    # loss_logger.add(f'loss_u_{epoch}_{batch_idx}', loss_dict['CCSSLCeLoss'].detach().cpu().numpy())
+    # loss_logger.add(f'loss_c_{epoch}_{batch_idx}', loss_dict['SoftSupConLoss'].detach().cpu().numpy())
+    # loss_logger.add(f'mask_prob_{epoch}_{batch_idx}', mask.mean().detach().numpy())
+    # out_logger.save('../align/step3/data/paddle_out.npy')
+    # loss_logger.save('../align/step3/data/paddle_loss.npy')
+    ##############################################################
+    # assert 1==0
+    return loss_dict, logits_x
+    
\ No newline at end of file

ppcls/loss/__init__.py

@@ -16,11 +16,13 @@ from .trihardloss import TriHardLoss
 from .triplet import TripletLoss, TripletLossV2
 from .tripletangularmarginloss import TripletAngularMarginLoss, TripletAngularMarginLoss_XBM
 from .supconloss import SupConLoss
+from .softsuploss import SoftSupConLoss
+from .ccssl_loss import CCSSLCeLoss
 from .pairwisecosface import PairwiseCosface
 from .dmlloss import DMLLoss
 from .distanceloss import DistanceLoss
 from .softtargetceloss import SoftTargetCrossEntropy
-
+from .ccssl_loss import CCSSLLoss
 from .distillationloss import DistillationCELoss
 from .distillationloss import DistillationGTCELoss
 from .distillationloss import DistillationDMLLoss

ppcls/loss/ccssl_loss.py

+
+from ppcls.engine.train.train import forward
+from .softsuploss import SoftSupConLoss
+import copy
+import paddle.nn as nn
+
+
+class CCSSLLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(CCSSLLoss, self).__init__()
+        ce_cfg = copy.deepcopy(kwargs['CELoss'])
+        self.ce_weight = ce_cfg.pop('weight')
+        softsupconloss_cfg = copy.deepcopy(kwargs['SoftSupConLoss'])
+        self.softsupconloss_weight = softsupconloss_cfg.pop('weight')
+
+        self.softsuploss = SoftSupConLoss(**softsupconloss_cfg)
+        self.celoss = nn.CrossEntropyLoss(reduction='none')
+
+    def forward(self, feats, batch, **kwargs):
+        """
+        Args:
+            feats: feature of s1 and s2, (n, 2, d)
+            batch: dict 
+        """
+        logits_w = batch['logits_w']
+        logits_s1 = batch['logits_s1']
+        p_targets_u_w = batch['p_targets_u_w']
+        mask = batch['mask']
+
+        max_probs = batch['max_probs']
+        # reduction = batch['reduction']
+       
+
+        loss_u = self.celoss(logits_s1, p_targets_u_w) * mask
+        loss_u = loss_u.mean()
+
+        loss_c = self.softsuploss(feats, max_probs, p_targets_u_w)
+
+        return {'CCSSLLoss': self.ce_weight*loss_u + self.softsupconloss_weight * loss_c}
+
+
+
+
+
+class CCSSLCeLoss(nn.Layer):
+    def __init__(self, **kwargs):
+        super(CCSSLCeLoss, self).__init__()
+        self.celoss = nn.CrossEntropyLoss(reduction='none')
+
+
+    def forward(self, inputs, batch, **kwargs):
+        p_targets_u_w = batch['p_targets_u_w']
+        logits_s1 = batch['logits_s1']
+        mask = batch['mask']
+        loss_u = self.celoss(logits_s1, p_targets_u_w) * mask
+        loss_u = loss_u.mean()
+
+        return {'CCSSLCeLoss': loss_u}
+
+
+

ppcls/loss/softsuploss.py

+"""
+CCSSL loss
+author: zhhike
+"""
+
+import paddle
+import paddle.nn as nn
+
+class SoftSupConLoss(nn.Layer):
+    """
+    Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
+    It also supports the unsupervised contrastive loss in SimCLR
+    """
+    def __init__(self, temperature=0.07, contrast_mode='all', base_temperature=0.07):
+        super(SoftSupConLoss, self).__init__()
+        self.temperature = temperature
+        self.contrast_mode = contrast_mode
+        self.base_temperature = base_temperature
+
+    def __call__(self, feat, batch, max_probs=None, labels=None, mask=None, reduction="mean", select_matrix=None):
+        """Compute loss for model. If both `labels` and `mask` are None,
+        it degenerates to SimCLR unsupervised loss:
+        https://arxiv.org/pdf/2002.05709.pdf
+
+        Args:
+            feat: hidden vector of shape [batch_size, n_views, ...].
+            labels: ground truth of shape [batch_size].
+            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
+                has the same class as sample i. Can be asymmetric.
+        Returns:
+            A loss scalar.
+        """
+        max_probs = batch['max_probs']
+        labels = batch['p_targets_u_w']
+        # reduction = batch['reduction']
+        batch_size = feat.shape[0]
+        if labels is not None:
+            labels = labels.reshape((-1, 1))
+            mask = paddle.equal(labels, labels.T).astype('float32')
+            max_probs = max_probs.reshape((-1, 1))
+            
+            score_mask = paddle.matmul(max_probs, max_probs.T)
+            
+            
+            mask = paddle.multiply(mask, score_mask)
+            
+        contrast_count = feat.shape[1]
+        contrast_feat = paddle.concat(paddle.unbind(feat, axis=1), axis=0)  # (2n, d)
+        if self.contrast_mode == 'all':
+            anchor_feat = contrast_feat
+            anchor_count = contrast_count
+        anchor_dot_contrast = paddle.matmul(anchor_feat, contrast_feat.T) / self.temperature
+        logits_max = anchor_dot_contrast.max(axis=1, keepdim=True)
+        logits = anchor_dot_contrast - logits_max.detach()
+        mask = paddle.concat([mask, mask], axis=0)
+        mask = paddle.concat([mask, mask], axis=1)
+        
+        # mask = paddle.repeat_interleave(paddle.repeat_interleave(mask, 2, 0), 2, 1)
+        logits_mask = 1 - paddle.eye(batch_size * contrast_count, dtype=paddle.float64)
+        mask = mask * logits_mask
+        exp_logits = paddle.exp(logits) * logits_mask
+        log_prob = logits - paddle.log(exp_logits.sum(axis=1, keepdim=True))
+        
+        mean_log_prob_pos = (mask * log_prob).sum(axis=1) / mask.sum(axis=1)
+        loss = -(self.temperature / self.base_temperature) * mean_log_prob_pos
+        loss = loss.reshape((anchor_count, batch_size))
+        if reduction == 'mean':
+            loss = loss.mean()
+
+        return {"SoftSupConLoss": loss}
+    
+
+

ppcls/optimizer/learning_rate.py

@@ -519,3 +519,21 @@ class CosineFixmatch(LRBase):
             last_epoch=self.last_epoch)
         setattr(learning_rate, "by_epoch", self.by_epoch)
         return learning_rate
+
+
+
+def cosine_schedule_with_warmup(learning_rate,
+                                num_warmup_steps,
+                                num_training_steps,
+                                num_cycles=7./16,
+                                last_epoch=-1,
+                                **kwargs):
+    
+    def _lr_lambda(current_step):
+        if current_step < num_warmup_steps:
+            return float(current_step) /  float(max(1, num_warmup_steps))
+        no_progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return max(0., math.cos(math.pi * num_cycles * no_progress))
+
+    return lr.LambdaDecay(learning_rate=learning_rate, lr_lambda=_lr_lambda)
+    
\ No newline at end of file