记一次图像分类排错

[rank6]: UnboundLocalError: local variable 'avg' referenced before assignment
File "train.py", line 289, in train_one_epoch
    logger.info(f"Rank {rank} epoch {epoch} avg acc: {avg_acc}")

def train_one_epoch(model, dataloader, criterion, optimizer, rank):
    losses = []
    correct = 0
    total = 0
    # 训练循环：若dataloader为空，这个循环完全不执行
    for images, labels in dataloader:
        # 前向传播（图像分类核心逻辑）
        outputs = model(images)
        loss = criterion(outputs, labels)
        losses.append(loss.item())
        
        # 计算准确率
        _, preds = torch.max(outputs, 1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    # 问题1：若循环没执行，sum(losses)/len(losses)不会执行，avg_loss未赋值
    avg_loss = sum(losses) / len(losses)
    # 问题2：若total=0，correct/total会除以0，且avg_acc未赋值
    avg_acc = correct / total
    
    # 引用未赋值的avg_loss/avg_acc，rank6无数据时直接报错
    print(f"Rank {rank} avg loss: {avg_loss}, avg acc: {avg_acc}")
    return avg_loss, avg_acc

def train_one_epoch(model, dataloader, criterion, optimizer, rank):
    losses = []
    correct = 0
    total = 0
    # 修复1：初始化avg变量，避免未赋值（关键！）
    avg_loss = 0.0
    avg_acc = 0.0
    
    for images, labels in dataloader:
        outputs = model(images)
        loss = criterion(outputs, labels)
        losses.append(loss.item())
        
        _, preds = torch.max(outputs, 1)
        correct += (preds == labels).sum().item()
        total += labels.size(0)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    # 修复2：加条件判断，仅当有数据时计算平均值
    if len(losses) > 0:
        avg_loss = sum(losses) / len(losses)
    else:
        # 空数据时记录日志，方便排查
        print(f"[Rank {rank}] Warning: 无训练数据，avg_loss设为0.0")
    
    # 修复3：避免除以0，同时处理空数据
    if total > 0:
        avg_acc = correct / total
    else:
        print(f"[Rank {rank}] Warning: 无有效样本，avg_acc设为0.0")
    
    # 此时avg_loss/avg_acc必有值，不会报错
    print(f"Rank {rank} avg loss: {avg_loss}, avg acc: {avg_acc}")
    return avg_loss, avg_acc

from torch.utils.data.distributed import DistributedSampler

# 构建分布式采样器（图像分类训练的标准配置）
train_sampler = DistributedSampler(
    train_dataset,  # 你的图像分类训练数据集
    num_replicas=world_size,  # rank总数（比如8）
    rank=rank,  # 当前rank
    drop_last=True  # 核心！丢弃最后一个不完整批次，避免某rank无数据
)

# 构建dataloader
train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,  # 批次大小
    sampler=train_sampler,  # 使用分布式采样器
    num_workers=4,  # 数据加载线程
    pin_memory=True  # 加速GPU数据传输
)