SuObjectDetection/train.py at main · Core43/SuObjectDetection · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import json
import torch
from torch.utils.data import DataLoader
import model
from torchvision import transforms
import data
import utils
def train(model, dataset, batch_size=16, epcho=50, lr=1e-3, device="cuda:0"):
    log = []
    t_p = transforms.ToPILImage()
    t_t = transforms.ToTensor()
    model = model.to(device)
    # 损失函数使用交叉熵损失函数
    lossfun = utils.VOLoss().to(device)
    # 选用Adam优化器，只在优化器中放入requires_grad为True的参数
    optim = torch.optim.Adam(model.parameters(), lr=lr)
    # dataloader
    dataloader = DataLoader(dataset, shuffle=True, batch_size=batch_size, drop_last=True)
    for i in range(epcho):
        # 损失和
        sum_loss = 0
        sum_loss_x_y_h_w, sum_loss_p, sum_loss_c = 0, 0, 0
        all_batch = 0
        # 遍历
        for index, (x, label) in enumerate(dataloader):
            x = x.to(device)
            label = label.to(device)
            # 模型输出
            output = model(x)
            # 计算loss
            loss_x_y_h_w, loss_p, loss_c = lossfun(output, label)
            loss = loss_x_y_h_w + loss_p + loss_c
            optim.zero_grad()
            loss.backward()
            optim.step()
            # 总步数增加
            all_batch += 1
            # loss和
            sum_loss += loss.item()
            sum_loss_x_y_h_w += loss_x_y_h_w.item()
            sum_loss_p += loss_p.item()
            sum_loss_c += loss_c.item()
            if (i + index) % (i + 60) == 0:
                image = t_p(x[0].to("cpu"))
                x_box = output[0].to("cpu")
                t_box = label[0].to("cpu")
                i_a = utils.draw_rectangle(image, utils.get_bbox(t_box))
                i_b = utils.draw_rectangle(image, utils.NMS(utils.get_bbox(x_box)))
                t_p(torch.cat([t_t(i_a), t_t(i_b), t_t(image)], dim=-1)).save(f"log/e{i}i{index}.jpg")
            # 打印进度条
            print(
                f"\repcho:{i} index:{index} avg_loss:{round(sum_loss / all_batch, 6)} bbox:{round(sum_loss_x_y_h_w / all_batch, 6)} p:{round(sum_loss_p / all_batch, 6)} c:{round(sum_loss_c / all_batch, 6)} loss:{round(loss.item(), 6)}",
                end="")
        # 保存模型
        torch.save(model.state_dict(), "last.pth")
        # 保存loss相关值
        log.append({"epcho": i, "avg_loss": sum_loss / all_batch, "bbox_loss": sum_loss_x_y_h_w/ all_batch,
                    "p_loss": (sum_loss_p / all_batch), "c_loss": (sum_loss_c / all_batch)})
        with open("log.json", "w", encoding="utf-8") as fp:
            json.dump(log, fp)
        print(
            f"\repcho:{i} avg_loss:{sum_loss / all_batch} bbox:{sum_loss_x_y_h_w / all_batch} p:{sum_loss_p / all_batch} c:{sum_loss_c / all_batch}")


if __name__ == '__main__':
    md = model.VRes(3, 5 + 4)
    print(sum([i.nelement() for i in md.parameters()]))
    dataset = data.ODData(r"E:\pythonPro\deep_learning\CV_objDetection_0\data", num_classes=4)
    train(md, dataset)