Resnet50進行遷移學習實現圖片二分類
內容簡介
本文使用預訓練的Resnet50網絡對皮膚病圖片進行二分類，基於portch框架。
數據集說明
數據集存放目錄爲： used_dataset , 共200張圖片，標籤爲：benign(良性）、malignant(患病)。
- 數據集劃分如下：
| 數據集類型 | benign | malignant | totall |
|---|---|---|---|
| train | 64 | 64 | 128 |
| val | 16 | 16 | 32 |
| test | 20 | 20 | 40 |
代碼目錄介紹
- args.py 存放訓練和測試所用的各種參數。
–mode字段表示運行模式：train or test.
–model_path字段是訓練模型的保存路徑。
其餘字段都有默認值。 - create_dataset.py 該腳本是用來讀json中的數據的，可以忽略。
- data_gen.py 該腳本實現劃分數據集以及數據增強和數據加載。
- main.py 包含訓練、評估和測試。
- transform.py 實現圖片增強。
- utils.py 存放一些工具函數。
- models/Res.py 是重寫的ResNet各種類型的網絡。
- checkpoints 保存模型
運行命令
# 訓練模型
python main.py --mode=train
# 測試模型
python main.py --mode=test --model_path='訓練好的模型文件路徑'
main.py 腳本介紹
main()函數 實現模型的訓練和評估
- step1: 加載數據
# data
transformations = get_transforms(input_size=args.image_size,test_size=args.image_size)
train_set = data_gen.Dataset(root=args.train_txt_path,transform=transformations['val_train'])
train_loader = data.DataLoader(train_set,batch_size=args.batch_size,shuffle=True)
val_set = data_gen.ValDataset(root=args.val_txt_path,transform=transformations['val_test'])
val_loader = data.DataLoader(val_set,batch_size=args.batch_size,shuffle=False)
- step2: 構建模型
model = make_model(args)
if use_cuda:
model.cuda()
# define loss function and optimizer
if use_cuda:
criterion = nn.CrossEntropyLoss().cuda()
else:
criterion = nn.CrossEntropyLoss()
optimizer = get_optimizer(model,args)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, patience=5, verbose=False)
- step3: 模型的訓練和評估
# train
for epoch in range(start_epoch,args.epochs):
print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, optimizer.param_groups[0]['lr']))
train_loss,train_acc = train(train_loader,model,criterion,optimizer,epoch,use_cuda)
test_loss,val_acc = val(val_loader,model,criterion,epoch,use_cuda)
scheduler.step(test_loss)
print(f'train_loss:{train_loss}\t val_loss:{test_loss}\t train_acc:{train_acc} \t val_acc:{val_acc}')
# save_model
is_best = val_acc >= best_acc
best_acc = max(val_acc, best_acc)
save_checkpoint({
'fold': 0,
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'train_acc':train_acc,
'acc': val_acc,
'best_acc': best_acc,
'optimizer' : optimizer.state_dict(),
}, is_best, single=True, checkpoint=args.checkpoint)
print("best acc = ",best_acc)
train()函數 每個epoch下的模型訓練過程
- 主要實現每個批次下梯度的反向傳播，計算accuarcy 和 loss, 並更新，最後返回其均值。
def train(train_loader,model,criterion,optimizer,epoch,use_cuda):
model.train()
losses = AverageMeter()
train_acc = AverageMeter()
for (inputs,targets) in tqdm(train_loader):
if use_cuda:
inputs,targets = inputs.cuda(),targets.cuda(async=True)
inputs,targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets)
# 梯度參數設爲0
optimizer.zero_grad()
# forward
outputs = model(inputs)
loss = criterion(outputs, targets)
# compute gradient and do SGD step
loss.backward()
optimizer.step()
# measure accuracy and record loss
acc = accuracy(outputs.data,targets.data)
# inputs.size(0)=32
losses.update(loss.item(), inputs.size(0))
train_acc.update(acc.item(),inputs.size(0))
return losses.avg,train_acc.avg
val()函數 每個epoch下的模型評估過程
- 主要代碼與train()函數一致，但沒有梯度的計算，還有將model.train()改成model.eval()。
def val(val_loader,model,criterion,epoch,use_cuda):
global best_acc
losses = AverageMeter()
val_acc = AverageMeter()
model.eval() # 將模型設置爲驗證模式
# 混淆矩陣
confusion_matrix = meter.ConfusionMeter(args.num_classes)
for _,(inputs,targets) in enumerate(val_loader):
if use_cuda:
inputs,targets = inputs.cuda(),targets.cuda()
inputs,targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets)
# compute output
outputs = model(inputs)
loss = criterion(outputs, targets)
confusion_matrix.add(outputs.data.squeeze(),targets.long())
acc1 = accuracy(outputs.data,targets.data)
# measure accuracy and record loss
losses.update(loss.item(), inputs.size(0))
val_acc.update(acc1.item(),inputs.size(0))
return losses.avg,val_acc.avg
test()函數 模型的測試
def test(use_cuda):
# data
transformations = get_transforms(input_size=args.image_size,test_size=args.image_size)
test_set = data_gen.TestDataset(root=args.test_txt_path,transform= transformations['test'])
test_loader = data.DataLoader(test_set,batch_size=args.batch_size,shuffle=False)
# load model
model = make_model(args)
if args.model_path:
# 加載模型
model.load_state_dict(torch.load(args.model_path))
if use_cuda:
model.cuda()
# evaluate
y_pred = []
y_true = []
img_paths = []
with torch.no_grad():
model.eval() # 設置成eval模式
for (inputs,targets,paths) in tqdm(test_loader):
y_true.extend(targets.detach().tolist())
img_paths.extend(list(paths))
if use_cuda:
inputs,targets = inputs.cuda(),targets.cuda()
inputs,targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets)
# compute output
outputs = model(inputs) # (16,2)
# dim=1 表示按行計算 即對每一行進行softmax
# probability = torch.nn.functional.softmax(outputs,dim=1)[:,1].tolist()
# probability = [1 if prob >= 0.5 else 0 for prob in probability]
# 返回最大值的索引
probability = torch.max(outputs, dim=1)[1].data.cpu().numpy().squeeze()
y_pred.extend(probability)
print("y_pred=",y_pred)
accuracy = metrics.accuracy_score(y_true,y_pred)
print("accuracy=",accuracy)
confusion_matrix = metrics.confusion_matrix(y_true,y_pred)
print("confusion_matrix=",confusion_matrix)
print(metrics.classification_report(y_true,y_pred))
# fpr,tpr,thresholds = metrics.roc_curve(y_true,y_pred)
print("roc-auc score=",metrics.roc_auc_score(y_true,y_pred))
res_dict = {
'img_path':img_paths,
'label':y_true,
'predict':y_pred,
}
df = pd.DataFrame(res_dict)
df.to_csv(args.result_csv,index=False)
print(f"write to {args.result_csv} succeed ")
實驗結果
