USER
import os
import pandas as pd
from PIL import Image
from tqdm import tqdm
import numpy as np
from sklearn.model_selection import train_test_split
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from torch.cuda.amp import GradScaler, autocast
from torch.optim.lr_scheduler import OneCycleLR
# 定义混合裁剪(Cutmix)的边界框函数
def rand_bbox(size, lam):
W, H = size[2], size[1]
cut_rat = np.sqrt(1. - lam)
cut_w = int(W * cut_rat)
cut_h = int(H * cut_rat)
cx = np.random.randint(W)
cy = np.random.randint(H)
bbx1 = np.clip(cx - cut_w // 2, 0, W)
bby1 = np.clip(cy - cut_h // 2, 0, H)
bbx2 = np.clip(cx + cut_w // 2, 0, W)
bby2 = np.clip(cy + cut_h // 2, 0, H)
return bbx1, bby1, bbx2, bby2
# 定义 DogBreedDataset 类
class DogBreedDataset(Dataset):
def __init__(self, csv_file, root_dir, transform=None, mixup=False, cutmix=False, alpha=1.0, indices=None):
self.labels_df = pd.read_csv(csv_file)
if indices is not None:
self.labels_df = self.labels_df.iloc[indices].reset_index(drop=True)
self.root_dir = root_dir
self.transform = transform
self.breeds = self.labels_df['breed'].unique()
self.breed_to_idx = {breed: idx for idx, breed in enumerate(self.breeds)}
self.labels_df['breed'] = self.labels_df['breed'].map(self.breed_to_idx)
self.mixup = mixup
self.cutmix = cutmix
self.alpha = alpha
def __len__(self):
return len(self.labels_df)
def __getitem__(self, idx):
img_id = self.labels_df.iloc[idx, 0]
breed = self.labels_df.iloc[idx, 1]
img_name = os.path.join(self.root_dir, img_id + '.jpg')
image = Image.open(img_name).convert('RGB')
if self.transform:
image = self.transform(image)
if self.mixup or self.cutmix:
rand_idx = np.random.randint(0, len(self))
img_id2 = self.labels_df.iloc[rand_idx, 0]
breed2 = self.labels_df.iloc[rand_idx, 1]
img_name2 = os.path.join(self.root_dir, img_id2 + '.jpg')
image2 = Image.open(img_name2).convert('RGB')
if self.transform:
image2 = self.transform(image2)
if self.mixup and np.random.rand() < 0.5:
lam = np.random.beta(self.alpha, self.alpha)
image = lam * image + (1 - lam) * image2
lam = float(lam) # 转换为浮点数
return image, breed, breed2, lam
elif self.cutmix:
lam = np.random.beta(self.alpha, self.alpha)
bbx1, bby1, bbx2, bby2 = rand_bbox(image.size(), lam)
image[:, bbx1:bbx2, bby1:bby2] = image2[:, bbx1:bbx2, bby1:bby2]
lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (image.size(-1) * image.size(-2)))
lam = float(lam) # 转换为浮点数
return image, breed, breed2, lam
return image, breed, breed, 1.0 # 返回相同数量的元素,但不应用混合
def get_breed_name(self, idx):
return self.breeds[idx]
# 定义 TestDataset 类
class TestDataset(Dataset):
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.image_ids = sorted(os.listdir(root_dir))
def __len__(self):
return len(self.image_ids)
def __getitem__(self, idx):
img_id = self.image_ids[idx].split('.')[0]
img_name = os.path.join(self.root_dir, self.image_ids[idx])
image = Image.open(img_name).convert('RGB')
if self.transform:
image = self.transform(image)
return image, img_id
# 数据转换
data_transform = transforms.Compose([
transforms.RandomResizedCrop(224, scale=(0.08, 1.0), ratio=(3.0/4.0, 4.0/3.0)),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),
transforms.RandomRotation(15),
transforms.RandomAffine(degrees=0, translate=(0.1, 0.1), scale=(0.9, 1.1)),
transforms.GaussianBlur(kernel_size=(5, 7), sigma=(0.1, 1.5)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
val_transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# 定义 ResNet152 模型
def get_model(num_classes):
model = models.resnet152(weights=models.ResNet152_Weights.IMAGENET1K_V2)
for param in model.parameters():
param.requires_grad = True
num_ftrs = model.fc.in_features
model.fc = nn.Sequential(
nn.Dropout(0.5),
nn.Linear(num_ftrs, 2048),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(2048, num_classes)
)
return model
# 训练函数
def train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device, grad_clip=1.0):
scaler = GradScaler()
best_val_loss = float('inf')
patience = 10
counter = 0
for epoch in range(num_epochs):
model.train()
train_loss = 0.0
train_correct = 0
train_total = 0
pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs} - Training")
for inputs, targets1, targets2, lam in pbar:
inputs = inputs.to(device)
targets1 = targets1.to(device)
targets2 = targets2.to(device)
lam = torch.tensor(lam, dtype=torch.float32).to(device) # 转换为张量并移动到设备
optimizer.zero_grad()
with autocast():
outputs = model(inputs)
loss1 = criterion(outputs, targets1) # shape: [batch_size]
loss2 = criterion(outputs, targets2) # shape: [batch_size]
loss = lam * loss1 + (1 - lam) * loss2 # shape: [batch_size]
loss = loss.mean()
if torch.isnan(loss):
print(f"发现NaN损失值在第{epoch+1}轮,停止训练。")
return best_val_loss
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
scaler.step(optimizer)
scaler.update()
if scheduler is not None:
scheduler.step()
train_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
# 对于混合标签,准确率的计算需要特别处理,这里暂且使用targets1作为参考
train_total += targets1.size(0)
train_correct += predicted.eq(targets1).sum().item()
pbar.set_postfix({
'loss': f"{loss.item():.4f}",
'acc': f"{train_correct/train_total:.4f}",
'lr': f"{optimizer.param_groups[0]['lr']:.6f}"
})
train_loss /= len(train_loader.dataset)
train_acc = train_correct / train_total
val_loss, val_acc = validate_model(model, val_loader, criterion, device)
print(f'\n第{epoch+1}轮/{num_epochs}轮:')
print(f'训练损失: {train_loss:.4f} 准确率: {train_acc:.4f}')
print(f'验证损失: {val_loss:.4f} 准确率: {val_acc:.4f}')
print(f'学习率: {optimizer.param_groups[0]["lr"]:.6f}')
if val_loss < best_val_loss:
best_val_loss = val_loss
torch.save(model.state_dict(), 'best_model.pth')
print("保存最佳模型。")
counter = 0
else:
counter += 1
if counter >= patience:
print(f"在第{epoch+1}轮达到早停。")
break
return best_val_loss
def validate_model(model, val_loader, criterion, device):
model.eval()
val_loss = 0.0
val_correct = 0
val_total = 0
with torch.no_grad():
for inputs, labels, _, _ in tqdm(val_loader, desc="验证中"):
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item() * inputs.size(0)
_, predicted = outputs.max(1)
val_total += labels.size(0)
val_correct += predicted.eq(labels).sum().item()
val_loss /= len(val_loader.dataset)
val_acc = val_correct / val_total
return val_loss, val_acc
# 测试函数
def test(model, test_loader, device, class_names):
predictions = []
model.eval()
with torch.no_grad():
for inputs, image_ids in tqdm(test_loader, desc="测试中"):
inputs = inputs.to(device)
outputs = model(inputs)
probs = torch.softmax(outputs, dim=1)
probs = probs.cpu().numpy()
for image_id, prob in zip(image_ids, probs):
prediction = {'id': image_id}
for i, class_name in enumerate(class_names):
prediction[class_name] = prob[i]
predictions.append(prediction)
return predictions
# 主训练循环
def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
labels_csv = '/kaggle/input/dog-breed-identification/labels.csv'
root_dir = '/kaggle/input/dog-breed-identification/train'
labels_df = pd.read_csv(labels_csv)
num_classes = len(labels_df['breed'].unique())
num_epochs = 100
batch_size = 32
dataset_size = len(labels_df)
indices = list(range(dataset_size))
train_indices, val_indices = train_test_split(
indices, test_size=0.2, random_state=42, stratify=labels_df['breed']
)
train_dataset = DogBreedDataset(
csv_file=labels_csv,
root_dir=root_dir,
transform=data_transform,
mixup=True,
cutmix=True,
alpha=1.0,
indices=train_indices
)
val_dataset = DogBreedDataset(
csv_file=labels_csv,
root_dir=root_dir,
transform=val_transform,
mixup=False,
cutmix=False,
alpha=1.0,
indices=val_indices
)
train_loader = DataLoader(
train_dataset, batch_size=batch_size, shuffle=True,
num_workers=4, pin_memory=True
)
val_loader = DataLoader(
val_dataset, batch_size=batch_size, shuffle=False,
num_workers=4, pin_memory=True
)
model = get_model(num_classes).to(device)
# 调整损失函数,设置 reduction='none' 以便进行加权损失计算
criterion = nn.CrossEntropyLoss(label_smoothing=0.1, reduction='none')
# 简化优化器,统一使用一个学习率
optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
scheduler = OneCycleLR(
optimizer,
max_lr=1e-3,
steps_per_epoch=len(train_loader),
epochs=num_epochs,
pct_start=0.3, # 预热阶段占比
anneal_strategy='cos', # 使用余弦退火策略
div_factor=10, # 初始学习率 / 最大学习率的比例
final_div_factor=1e4 # 最大学习率 / 最终学习率的比例
)
best_loss = train_model(
model, train_loader, val_loader, criterion, optimizer,
scheduler, num_epochs, device, grad_clip=1.0
)
print(f"最佳验证损失: {best_loss:.4f}")
best_model = get_model(num_classes).to(device)
best_model.load_state_dict(torch.load('best_model.pth'))
return best_model, labels_df
if __name__ == "__main__":
best_model, labels_df = main()
# 准备测试数据
test_dataset = TestDataset(
root_dir='/kaggle/input/dog-breed-identification/test',
transform=val_transform
)
test_loader = DataLoader(
test_dataset, batch_size=64, shuffle=False,
num_workers=4, pin_memory=True
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 获取所有类别名称
class_names = labels_df['breed'].unique().tolist()
# 进行测试
predictions = test(best_model, test_loader, device, class_names)
# 创建 DataFrame 并保存为 CSV
submission_df = pd.DataFrame(predictions)
submission_df.to_csv('/kaggle/working/submission.csv', index=False)
print("预测结果已保存到 submission.csv")
# 添加一些额外的打印语句来帮助调试
print("训练完成")
print(f"设备: {device}")
print(f"类别数量: {len(class_names)}")
print(f"测试集大小: {len(test_dataset)}")
报错
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
Cell In[4], line 328
325 return best_model, labels_df
327 if __name__ == "__main__":
--> 328 best_model, labels_df = main()
330 # 准备测试数据
331 test_dataset = TestDataset(
332 root_dir='/kaggle/input/dog-breed-identification/test',
333 transform=val_transform
334 )
Cell In[4], line 316, in main()
303 optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
305 scheduler = OneCycleLR(
306 optimizer,
307 max_lr=1e-3,
(...)
313 final_div_factor=1e4 # 最大学习率 / 最终学习率的比例
314 )
--> 316 best_loss = train_model(
317 model, train_loader, val_loader, criterion, optimizer,
318 scheduler, num_epochs, device, grad_clip=1.0
319 )
321 print(f"最佳验证损失: {best_loss:.4f}")
323 best_model = get_model(num_classes).to(device)
Cell In[4], line 194, in train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs, device, grad_clip)
191 train_loss /= len(train_loader.dataset)
192 train_acc = train_correct / train_total
--> 194 val_loss, val_acc = validate_model(model, val_loader, criterion, device)
196 print(f'\n第{epoch+1}轮/{num_epochs}轮:')
197 print(f'训练损失: {train_loss:.4f} 准确率: {train_acc:.4f}')
Cell In[4], line 225, in validate_model(model, val_loader, criterion, device)
222 outputs = model(inputs)
223 loss = criterion(outputs, labels)
--> 225 val_loss += loss.item() * inputs.size(0)
226 _, predicted = outputs.max(1)
227 val_total += labels.size(0)
RuntimeError: a Tensor with 32 elements cannot be converted to Scalar
如何解决