167 lines
5.9 KiB
Python
Executable File
167 lines
5.9 KiB
Python
Executable File
'''
|
||
author ： Hatcher W
|
||
Email： hcwang@semi.ac.cn
|
||
use: 密度估计训练
|
||
'''
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch import optim
|
||
from torch.utils.data import DataLoader
|
||
import numpy as np
|
||
from network import CSRNet,PKSNet,PKSNet_attention
|
||
from dataset import People_Dataset
|
||
from configs.config import *
|
||
from configs.config import model_head
|
||
import json
|
||
import os
|
||
from tqdm import tqdm
|
||
import cv2
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
|
||
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
||
train_weigth_path = "../data_MPKFNet/weigth/BirdCounting2/"
|
||
if not os.path.exists(train_weigth_path):
|
||
os.makedirs(train_weigth_path)
|
||
'''
|
||
准备数据
|
||
'''
|
||
batch_size = 1
|
||
with open(trainA_json, 'r') as outfile:
|
||
train_list = json.load(outfile)
|
||
with open(testA_json, 'r') as outfile:
|
||
test_list = json.load(outfile)
|
||
train_dataset = People_Dataset(train_list)
|
||
train_loader = DataLoader(train_dataset, batch_size=batch_size)
|
||
test_dataset = People_Dataset(test_list)
|
||
test_loader = DataLoader(test_dataset, batch_size=1)
|
||
|
||
'''
|
||
创建网络
|
||
'''
|
||
net = PKSNet().to(device)
|
||
|
||
|
||
'''
|
||
训练设置
|
||
'''
|
||
lr=1e-10
|
||
criterion_desity = nn.MSELoss(reduction='sum')
|
||
criterion_spice = nn.CrossEntropyLoss()
|
||
criterion_posture = nn.CrossEntropyLoss()
|
||
optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.95, weight_decay=5*1e-4)
|
||
epoch_start = 0
|
||
# optimizer = optim.Adam(net.parameters(), lr=1e-7, weight_decay=1e-4)
|
||
if True:
|
||
# net.load_state_dict(torch.load('./weigth/bestmae.pth'))
|
||
checkpoint = torch.load(train_weigth_path+'final.pth')
|
||
net.load_state_dict(checkpoint['net'])
|
||
#optimizer.load_state_dict(checkpoint['optimizer'])
|
||
#epoch_start = checkpoint['epoch']
|
||
print("load model weight ok")
|
||
|
||
'''
|
||
开始训练
|
||
'''
|
||
mae_best = 1e10
|
||
num_lr=0
|
||
print(net)
|
||
for epoch in range(epoch_start,epochs):
|
||
net.train()
|
||
epoch_loss = 0
|
||
with tqdm(total=len(train_dataset), desc=f'Epoch {epoch + 1}/{epochs}', unit='img') as pbar:
|
||
if model_head:
|
||
for img, target, count, posture, species in train_loader:
|
||
img = img.to(device)
|
||
target = target.to(device)
|
||
posture = posture.to(device)
|
||
species = species.to(device)
|
||
target_pred, posture_pred, spice_pred = net(img)
|
||
# print(target_pred.shape, posture_pred.shape, spice_pred.shape)
|
||
# print(target.shape, posture.shape, species.shape)
|
||
loss = 0.9*criterion_desity(target_pred, target) + 0.05*criterion_posture(posture_pred, posture) \
|
||
+ 0.05*criterion_spice(spice_pred, species)
|
||
epoch_loss += loss.item()
|
||
pbar.set_postfix({'loss (batch)': loss.item()/img.shape[0]})
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
optimizer.step()
|
||
pbar.update(img.shape[0])
|
||
else:
|
||
for img, target, count in train_loader:
|
||
img = img.to(device)
|
||
target = target.to(device)
|
||
target_pred= net(img)
|
||
# print(target_pred.shape, posture_pred.shape, spice_pred.shape)
|
||
# print(target.shape, posture.shape, species.shape)
|
||
loss = criterion_desity(target_pred, target)
|
||
epoch_loss += loss.item()
|
||
pbar.set_postfix({'loss (batch)': loss.item() / img.shape[0]})
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
optimizer.step()
|
||
pbar.update(img.shape[0])
|
||
print('train_data: {}, iter loss: {} for single image \n'.format(len(train_list), epoch_loss/len(train_list)))
|
||
|
||
mae ,mse= 0,0
|
||
if model_head:
|
||
for img, traget, count, posture, species in tqdm(test_loader):
|
||
img = img.to(device)
|
||
target = traget.to(device)
|
||
targets_pred, posture_pred, spice_pred = net(img)
|
||
#print(targets_pred.data.sum())
|
||
#print(target.data.sum())
|
||
mae += abs(targets_pred.data.sum() - target.data.sum())
|
||
mse += ((target.data.sum() - targets_pred.data.sum()) * (target.data.sum() - targets_pred.data.sum()))
|
||
else:
|
||
for img, traget, count in tqdm(test_loader):
|
||
img = img.to(device)
|
||
target = traget.to(device)
|
||
targets_pred = net(img)
|
||
#print(targets_pred.data.sum())
|
||
#print(target.data.sum())
|
||
mae += abs(targets_pred.data.sum() - target.data.sum())
|
||
mse += ((target.data.sum() - targets_pred.data.sum()) * (target.data.sum() - targets_pred.data.sum()))
|
||
mae = mae / len(test_list)
|
||
mse = mse /len(test_list)
|
||
if mae < mae_best:
|
||
num_lr=0
|
||
print(epoch)
|
||
state = {
|
||
"epoch" :epoch,
|
||
"net" :net.state_dict(),
|
||
"optimizer":optimizer.state_dict()
|
||
|
||
}
|
||
torch.save(state, train_weigth_path+'bestmaesheep{},mse{}.pth'.format(int(mae),int(mse)))
|
||
torch.save(state, train_weigth_path+'bestmaesheep.pth'.format(int(mae)))
|
||
mae_best = mae
|
||
else:
|
||
num_lr+=1
|
||
if num_lr>5:
|
||
lr*=0.1
|
||
num_lr = 0
|
||
if lr < 1e-10:
|
||
lr = 1e-10
|
||
print("lr is down to :"+str(lr))
|
||
optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.95, weight_decay=5 * 1e-4)
|
||
print('test_data:{},calculate MAE: \n'.format(len(test_list)))
|
||
print(' * MAE {mae:.3f} \n'.format(mae=mae))
|
||
print(' * MSE {mse:.3f} \n'.format(mse=mse))
|
||
|
||
|
||
'''
|
||
模型加载与单张图片测试
|
||
'''
|
||
# net = CSRNet()
|
||
# net.load_state_dict(torch.load('CSRnet.pth'))
|
||
# net.eval()
|
||
# test_img = cv2.imread('./data/part_A_final/test_data/images/IMG_1.jpg.jpg')
|
||
# test_img = cv2.resize(test_img, (224, 224))
|
||
# test_img = np.transpose(test_img, [2, 0, 1])
|
||
# test_img = np.expand_dims(test_img,axis=0).astype(np.float)
|
||
# with torch.no_grad():
|
||
# out = net(torch.from_numpy(test_img).float())
|
||
# out = out.squeeze(0).numpy()
|
||
# cv2.imwrite('result.png',out)
|
||
|
||
|