对比学习用于分类任务

时间:2024-05-03 09:35:03

import os
os.environ["KERAS_BACKEND"] = "tensorflow"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import resource
low, high = resource.getrlimit(resource.RLIMIT_NOFILE)

resource.setrlimit(resource.RLIMIT_NOFILE, (high, high))

import math
import matplotlib.pyplot as plt
import tensorflow as tf
import keras
from keras import ops
from keras import layers

gpus = tf.config.list_physical_devices('GPU')  
if gpus:  
    # 如果有GPU,设置GPU资源使用率  
    try:  
        # 允许GPU内存按需增长  
        for gpu in gpus:  
            tf.config.experimental.set_memory_growth(gpu, True)  
        # 设置可见的GPU设备(这里实际上不需要,因为已经通过内存增长设置了每个GPU)  
        # tf.config.set_visible_devices(gpus, 'GPU')  
        print("GPU可用并已设置内存增长模式。")  
    except RuntimeError as e:  
        # 虚拟设备未就绪时可能无法设置GPU  
        print(f"设置GPU时发生错误: {e}")  
else:  
    # 如果没有GPU  
    print("没有检测到GPU设备。")

unlabeled_dataset_size = 50000
labeled_dataset_size = 5000
image_channels = 3
num_epochs = 30
batch_size =110
temperature = 0.1
contrastive_augmentation = {"min_area": 0.25, "brightness": 0.6, "jitter": 0.2}
classification_augmentation = {
    "min_area": 0.75,
    "brightness": 0.3,
    "jitter": 0.1,
}

import numpy as np

root_dir='./datasets/stl10_binary'

import random

def read_bin_image(filename, width, height):  
    # 打开二进制文件  
    with open(filename, 'rb') as f:  
        bin_data = f.read()
        #这里读取到的是每一个通道的像素数据
        pixels = np.frombuffer(bin_data, dtype=np.uint8)
        # print(len(pixels))
        # 重塑为图像的形状(n,c,w,h)-->(n,h,w,c),先reshape,是因为原来像素数据就是按那
        #种方式摊平的,之后换一下轴转换成tensorflow形式
        image = pixels.reshape(-1,3,width,height).transpose(0,3,2,1)  
    return image 

def read_bin_label(filename):
    with open(filename, 'rb') as f:
        labels=np.fromfile(f,dtype=np.uint8)
    return labels

filename=root_dir+'/train_X.bin'
x_train=read_bin_image(filename,96,96)

filename2=root_dir+'/train_y.bin'

y_train=read_bin_label(filename2)

y_train=y_train-1

classes=[line.strip()for line in open(root_dir+'/class_names.txt').readlines()]

plt.figure(figsize=(10,10))
x_t=random.sample(range(len(x_train)),9)#随机选取9个样本图片
for i in range(9):
    plt.subplot(3,3,i+1)
    plt.imshow(x_train[x_t[i]])#x_t[i]会依次获取其中的每一个索引
    plt.axis('off')
    # print(y_train[x_t[i]])
    plt.title(classes[int(y_train[x_t[i]])])
plt.show()

train_dataset=tf.data.Dataset.from_tensor_slices((x_train,y_train))

for i,j in train_dataset.take(1):
    print(i.shape,j,i.numpy().max(),i.numpy().min())
    plt.imshow(i)

filename=root_dir+'/test_X.bin'
x_test=read_bin_image(filename,96,96)

filename2=root_dir+'/test_y.bin'
y_test=read_bin_label(filename2)

y_test=y_test-1

plt.figure(figsize=(10,10))
x_t_2=random.sample(range(len(x_test)),9)#随机选取9个样本图片
for i in range(9):
    plt.subplot(3,3,i+1)
    plt.imshow(x_test[x_t_2[i]])#x_t[i]会依次获取其中的每一个索引
    plt.axis('off')
    # print(y_train[x_t[i]])
    plt.title(classes[int(y_test[x_t_2[i]])])
plt.show()

filename=root_dir+'/unlabeled_X.bin'
unlabed_x_train=read_bin_image(filename,96,96)

plt.figure(figsize=(10,10))
x_t_2=random.sample(range(len(u