【机器学习】--- 自监督学习

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.models as models
from torch.utils.data import DataLoader, Dataset
import numpy as np

# SimCLR数据增强
class SimCLRTransform:
    def __init__(self, size):
        self.transform = transforms.Compose([
            transforms.RandomResizedCrop(size=size),
            transforms.RandomHorizontalFlip(),
            transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
            transforms.RandomGrayscale(p=0.2),
            transforms.GaussianBlur(kernel_size=(3, 3)),
            transforms.ToTensor()
        ])

    def __call__(self, x):
        return self.transform(x), self.transform(x)

# 定义对比损失
class NTXentLoss(nn.Module):
    def __init__(self, temperature):
        super(NTXentLoss, self).__init__()
        self.temperature = temperature

    def forward(self, z_i, z_j):
        batch_size = z_i.size(0)
        z = torch.cat([z_i, z_j], dim=0)
        sim_matrix = torch.mm(z, z.t()) / self.temperature
        mask = torch.eye(2 * batch_size, dtype=torch.bool).to(sim_matrix.device)
        sim_matrix.masked_fill_(mask, -float('inf'))
        
        positives = torch.cat([torch.diag(sim_matrix, batch_size), torch.diag(sim_matrix, -batch_size)], dim=0)
        negatives = sim_matrix[~mask].view(2 * batch_size, -1)
        
        logits = torch.cat([positives.unsqueeze(1), negatives], dim=1)
        labels = torch.zeros(2 * batch_size).long().to(logits.device)
        
        loss = nn.CrossEntropyLoss()(logits, labels)
        return loss

# 定义模型架构
class SimCLR(nn.Module):
    def __init__(self, base_model, projection_dim=128):
        super(SimCLR, self).__init__()
        self.backbone = base_model
        self.projector = nn.Sequential(
            nn.Linear(self.backbone.fc.in_features, 512),
            nn.ReLU(),
            nn.Linear(512, projection_dim)
        )

    def forward(self, x):
        h = self.backbone(x)
        z = self.projector(h)
        return z

# 模型训练
def train_simclr(model, train_loader, epochs=100, lr=1e-3, temperature=0.5):
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = NTXentLoss(temperature)

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for x_i, x_j in train_loader:
            optimizer.zero_grad()
            z_i = model(x_i)
            z_j = model(x_j)
            loss = criterion(z_i, z_j)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {total_loss/len(train_loader)}')

# 示例：在CIFAR-10上进行SimCLR训练
from torchvision.datasets import CIFAR10

train_dataset = CIFAR10(root='./data', train=True, transform=SimCLRTransform(32), download=True)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)

resnet_model = models.resnet18(pretrained=False)
simclr_model = SimCLR(base_model=resnet_model)

train_simclr(simclr_model, train_loader)

import random

# 定义Jigsaw数据预处理
class JigsawTransform:
    def __init__(self, size, grid_size=3):
        self.size = size
        self.grid_size = grid_size
        self.transform = transforms.Compose([
            transforms.Resize((size, size)),
            transforms.ToTensor()
        ])

    def __call__(self, x):
        x = self.transform(x)
        blocks = self.split_into_blocks(x)
        random.shuffle(blocks)
        return torch.cat(blocks, dim=1), torch.tensor([i for i in range(self.grid_size ** 2)])

    def split_into_blocks(self, img):
        c, h, w = img.size()
        block_h, block_w = h // self.grid_size, w // self.grid_size
        blocks = []
        for i in range(self.grid_size):
            for j in range(self.grid_size):
                block = img[:, i*block_h:(i+1)*block_h, j*block_w:(j+1)*block_w]
                blocks.append(block.unsqueeze(0))
        return blocks

# 定义Jigsaw任务模型
class JigsawModel(nn.Module):
    def __init__(self, base_model):
        super(JigsawModel, self).__init__()
        self.backbone = base_model
        self.classifier = nn.Linear(base_model.fc.in_features, 9)

    def forward(self, x):
        features = self.backbone(x)
        out = self.classifier(features)
        return out

# 示例：在CIFAR-10上进行Jigsaw任务训练
train_dataset = CIFAR10(root='./data', train=True, transform=JigsawTransform(32), download=True)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)

jigsaw_model = JigsawModel(base_model=resnet_model)

# 训练过程同样可以采用类似SimCLR的方式进行