import torch
import torch.nn as nn
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os
import matplotlib.pyplot as plt
import torchvision.models as models
# This is for the progress bar.from tqdm import tqdm
import seaborn as sns
device = torch.device("cuda"if torch.cuda.is_available()else"cpu")# 是否要冻住模型的前面一些层defset_parameter_requires_grad(model, feature_extracting):# if feature_extracting:# model = model# for i, param in enumerate(model.children()):# if i == 8:# break# param.requires_grad = Falseif feature_extracting:
model = model
for param in model.parameters():
param.requires_grad =False# resnet34模型defres_model(num_classes, feature_extract =False, use_pretrained=True):
model_ft = models.resnet34(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
# model_ft.fc = nn.Sequential(# nn.Linear(num_ftrs, 512),# nn.ReLU(inplace=True),# nn.Dropout(.3),# nn.Linear(512, len(num_to_class))# )
model_ft.fc = nn.Sequential(
nn.Linear(num_ftrs, num_classes))return model_ft
# resnext50模型defresnext_model(num_classes, feature_extract =False, use_pretrained=True):
model_ft = models.resnext50_32x4d(pretrained=use_pretrained)
set_parameter_requires_grad(model_ft, feature_extract)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Sequential(nn.Linear(num_ftrs, num_classes))return model_ft
# 超参数设置
learning_rate =3e-4
weight_decay =1e-3
num_epoch =50
model_path =r'E:/Learn_MechineLearning_DeepLearning/pytorch-Learning/Kaggle/Classify-leaves/classify-leaves/pre_trained_model/pre_trained_resnext50.pthe'
Step 4:训练模型,保存,测试
# Initialize a model, and put it on the device specified.#model = res_model(176)
model = resnext_model(176)
model = model.to(device)
model.device = device
# For the classification task, we use cross-entropy as the measurement of performance.
criterion = nn.CrossEntropyLoss()# Initialize optimizer, you may fine-tune some hyperparameters such as learning rate on your own.# optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = learning_rate, weight_decay=weight_decay)
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate, weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=0, last_epoch=-1)# The number of training epochs.
n_epochs = num_epoch
best_acc =0.0for epoch inrange(n_epochs):# ---------- Training ----------# Make sure the model is in train mode before training.
model.train()# These are used to record information in training.
train_loss =[]
train_accs =[]
i =0# Iterate the training set by batches.for batch in tqdm(train_loader):# A batch consists of image data and corresponding labels.
imgs, labels = batch
imgs = imgs.to(device)
labels = labels.to(device)# Forward the data. (Make sure data and model are on the same device.)
logits = model(imgs)# Calculate the cross-entropy loss.# We don't need to apply softmax before computing cross-entropy as it is done automatically.
loss = criterion(logits, labels)# Gradients stored in the parameters in the previous step should be cleared out first.
optimizer.zero_grad()# Compute the gradients for parameters.
loss.backward()# Update the parameters with computed gradients.
optimizer.step()# 更新学习率------------------------------------------------------------------------------
scheduler.step()if(i %500==0):print("learning_rate:", scheduler.get_last_lr()[0])
i = i +1# Compute the accuracy for current batch.
acc =(logits.argmax(dim=-1)== labels).float().mean()# Record the loss and accuracy.
train_loss.append(loss.item())
train_accs.append(acc)# The average loss and accuracy of the training set is the average of the recorded values.
train_loss =sum(train_loss)/len(train_loss)
train_acc =sum(train_accs)/len(train_accs)# Print the information.print(f"[ Train | {epoch +1:03d}/{n_epochs:03d} ] loss = {train_loss:.5f}, acc = {train_acc:.5f}")# ---------- Validation ----------# Make sure the model is in eval mode so that some modules like dropout are disabled and work normally.
model.eval()# These are used to record information in validation.
valid_loss =[]
valid_accs =[]# Iterate the validation set by batches.for batch in tqdm(val_loader):
imgs, labels = batch
# We don't need gradient in validation.# Using torch.no_grad() accelerates the forward process.with torch.no_grad():
logits = model(imgs.to(device))# We can still compute the loss (but not the gradient).
loss = criterion(logits, labels.to(device))# Compute the accuracy for current batch.
acc =(logits.argmax(dim=-1)== labels.to(device)).float().mean()# Record the loss and accuracy.
valid_loss.append(loss.item())
valid_accs.append(acc)# The average loss and accuracy for entire validation set is the average of the recorded values.
valid_loss =sum(valid_loss)/len(valid_loss)
valid_acc =sum(valid_accs)/len(valid_accs)# Print the information.print(f"[ Valid | {epoch +1:03d}/{n_epochs:03d} ] loss = {valid_loss:.5f}, acc = {valid_acc:.5f}")# if the model improves, save a checkpoint at this epochif valid_acc > best_acc:
best_acc = valid_acc
torch.save(model.state_dict(), model_path)print('saving model with acc {:.3f}'.format(best_acc))# 测试集
saveFileName ='./submission.csv'## predict#model = res_model(176)
model = resnext_model(176)# create model and load weights from checkpoint
model = model.to(device)
model.load_state_dict(torch.load(model_path))# Make sure the model is in eval mode.# Some modules like Dropout or BatchNorm affect if the model is in training mode.
model.eval()# Initialize a list to store the predictions.
predictions =[]# Iterate the testing set by batches.for batch in tqdm(test_loader):
imgs = batch
with torch.no_grad():
logits = model(imgs.to(device))# Take the class with greatest logit as prediction and record it.
predictions.extend(logits.argmax(dim=-1).cpu().numpy().tolist())
preds =[]for i in predictions:
preds.append(num_to_class[i])
test_data = pd.read_csv(test_path)
test_data['label']= pd.Series(preds)
submission = pd.concat([test_data['image'], test_data['label']], axis=1)
submission.to_csv(saveFileName, index=False)print("Done!!!!!!!!!!!!!!!!!!!!!!!!!!!")