CS182_Project/train_3d_data.py at master · nagara214/CS182_Project · 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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from torch.utils.data import DataLoader
from torch.nn.parallel import DataParallel
import torch
import torch.nn as nn
from tqdm import tqdm
import os

from utils.utils import Config, min_max_normalization, setup_seeds, get_date_time
from utils.dataset import KaggleDataset
from models.unet import unet_build_model
from optimizer.loss import surface_dice

class DiceLoss(nn.Module):
    def __init__(self, smooth=0.00001):
        super(DiceLoss, self).__init__()
        self.smooth = smooth

    def forward(self, inputs, targets):
        inputs = torch.sigmoid(inputs)
        inputs = inputs.view(-1)
        targets = targets.view(-1)

        intersection = (inputs * targets).sum()
        dice = (2.*intersection + self.smooth)/(inputs.sum() + targets.sum() + self.smooth)

        return 1 - dice

def main():
    cfg = Config('configs/train_config.yaml')
    setup_seeds(cfg)
    train_dataset = KaggleDataset(cfg)
    train_loader = DataLoader(train_dataset, batch_size=16, num_workers=cfg.num_workers)
    val_dataset = KaggleDataset(cfg, mode='val')
    val_loader = DataLoader(val_dataset, batch_size=16, num_workers=cfg.num_workers)

    model = unet_build_model(cfg)
    model = DataParallel(model)

    loss_fn_1 = DiceLoss() # 和 BCEloss 加权 预训练 模型 医学 CT
    loss_fn_2 = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr)
    scaler = torch.cuda.amp.GradScaler()
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=cfg.lr,
                                                steps_per_epoch=len(train_loader), epochs=cfg.epochs+1,
                                                pct_start=0.1)

    date, time = get_date_time()

    for epoch in range(cfg.epochs):
        train_one_epoch(model, train_loader, loss_fn_1,loss_fn_2,  optimizer, scaler, scheduler, epoch, cfg)
        val(model, val_loader, loss_fn_1,loss_fn_2, epoch, cfg)
        torch.save(model.state_dict(), os.path.join(cfg.output_path, f'{cfg.model_name}_{date}_{time}_Pure_Dice/epoch_{epoch}.pt'))


def train_one_epoch(model, train_loader, loss_fn_1 ,loss_fn_2, optimizer, scaler, scheduler, epoch, cfg):
    model.train()
    iters = tqdm(range(len(train_loader)))
    train_loss = 0
    dice_score = 0

    for i, (images, masks) in enumerate(train_loader):
        images = images.cuda()
        masks = masks.cuda()
        images = min_max_normalization(images)

        with torch.cuda.amp.autocast():
            outputs = model(images)
            loss = loss_fn_2(outputs, masks)

        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()

        train_loss = (train_loss*i + loss.item())/(i+1)
        dice_score = (dice_score*i + surface_dice(outputs, masks))/(i+1)
        iters.set_description(f"Epoch {epoch+1}/{cfg.epochs}, train loss: {train_loss:.4f}, dice score: {dice_score:.4f}")
        iters.update()
    iters.close()


def val(model, val_loader, loss_fn_1,loss_fn_2, epoch, cfg):
    model.eval()
    iters = tqdm(range(len(val_loader)))
    val_loss = 0
    val_dice_score = 0

    for i, (images, masks) in enumerate(val_loader):
        images = images.cuda()
        masks = masks.cuda()
        images = min_max_normalization(images)

        with torch.cuda.amp.autocast():
            with torch.no_grad():
                outputs = model(images)
                loss = loss_fn_2(outputs, masks)

        val_loss = (val_loss*i + loss.item())/(i+1)
        val_dice_score = (val_dice_score*i + surface_dice(outputs, masks))/(i+1)
        iters.set_description(f"Epoch {epoch+1}/{cfg.epochs}, val loss: {val_loss:.4f}, val dice score: {val_dice_score:.4f}")
        iters.update()
    iters.close()


if __name__ == '__main__':
    main()