PyTorch清理CPU缓存的实现步骤

import torch
# 创建大量的Tensor对象
data = torch.randn(1000, 1000, 1000)
# 进行计算...
# 清理CPU缓存
torch.cuda.empty_cache()

import torch
from torch import nn, optim
from torchvision import models, datasets, transforms
# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 定义模型
model = models.resnet50(pretrained=True)
model.to(device)
# 加载数据集
train_dataset = datasets.ImageFolder("train_data_path", transform=transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
]))
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
# 训练模型
for epoch in range(10):
    running_loss = 0.0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        # 清理CPU缓存
        torch.cuda.empty_cache()
        running_loss += loss.item()
    # 打印每个epoch的损失
    epoch_loss = running_loss / len(train_loader)
    print(f"Epoch [{epoch+1}/10], Loss: {epoch_loss:.4f}")

import torch
from torch.utils.data import Dataset, DataLoader
# 自定义数据集类
class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels):
        self.texts = texts
        self.labels = labels
    def __len__(self):
        return len(self.texts)
    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        return text, label
# 自定义collate_fn函数
def collate_fn(batch):
    texts, labels = zip(*batch)
    return texts, labels
# 创建数据集
texts = [...]  # 文本数据
labels = [...]  # 标签
dataset = TextClassificationDataset(texts, labels)
# 创建数据加载器
batch_size = 32
dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn)
# 模型训练
for texts, labels in dataloader:
    texts = preprocess(texts)  # 数据预处理
    inputs = tokenize(texts)  # 文本分词等处理
    outputs = model(inputs)  # 模型推理
    loss = calculate_loss(outputs, labels)  # 计算损失
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()