first commit

crow_train.py

@@ -0,0 +1,166 @@
+'''
+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)
+
+