-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
87 lines (70 loc) · 2.88 KB
/
train.py
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
import torch
import time
from tqdm import tqdm
from utils import *
import wandb
import torch.distributed as dist
def reduce_tensor(tensor, n):
rt = tensor.clone()
dist.all_reduce(rt, op=dist.ReduceOp.SUM)
rt /= n
return rt
def do_train(args, model, optimizer, criterion, train_dl, valid_dl, scheduler):
if args.is_master :
wandb.init(name = f"{args.model}", project = "Dacon_Segmentation", reinit = True, entity = args.wandb_id, config = args)
print("Stat Train and Valid")
best_loss = 1
train_loss = 0
scaler = torch.cuda.amp.GradScaler(enabled = True)
for epoch in range(args.epochs):
train_dl.sampler.set_epoch(epoch)
if args.is_master:
print(f"Epoch : {epoch + 1}")
model.train()
for img, mask in tqdm(train_dl):
img, mask = img.to(args.device, dtype=torch.float), mask.to(args.device, dtype=torch.float)
optimizer.zero_grad()
epoch_loss =0
with torch.cuda.amp.autocast(enabled = True):
outputs= model(img)
loss = criterion(outputs, mask.unsqueeze(1))
train_loss += reduce_tensor(loss, args.world_size) if args.distributed else loss
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
torch.cuda.synchronize()
scheduler.step()
train_loss /= len(train_dl)
test_loss, threshold = 0, 0.35
print(f"Validation step, epoch: {epoch + 1}")
model.eval()
dice_scores=0
for img, mask in valid_dl: # # valid code
img, mask = img.to(args.device, dtype=torch.float), mask.to(args.device, dtype=torch.float)
with torch.no_grad():
outputs = model(img)
masks = torch.sigmoid(outputs).squeeze(1)
masks = (masks>0.35).float()
# masks = torch.sigmoid(outputs).cpu().numpy()
# masks = np.squeeze(masks, axis=1)
# masks = (masks>0.35).astype(np.uint8)
loss = criterion(outputs, mask.unsqueeze(1))
ds = dice_score_torch(masks, mask)
if args.distributed:
ds = reduce_tensor(ds, args.world_size)
loss = reduce_tensor(loss, args.world_size)
dice_scores+= ds
test_loss += loss
test_loss = test_loss/len(valid_dl)
test_dice= dice_scores/len(valid_dl)
if args.is_master:
print(f" Loss : {test_loss}, Dice : {test_dice}")
if args.save_model and best_loss > test_loss:
best_loss = test_loss
print("save", test_loss)
save_model(args, model)
wandb.log({
"train Loss" : train_loss,
"test Loss" : test_loss,
"test Dice" : test_dice
})