-
Notifications
You must be signed in to change notification settings - Fork 13
/
train.py
186 lines (165 loc) · 10.8 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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
from ast import arg
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '3'
import argparse
from pickle import FALSE, TRUE
from statistics import mode
from tkinter import image_names
import torch
import torchvision
from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
import torch.optim as optim
import numpy as np
import torch
from torch.utils.tensorboard import SummaryWriter
import time
import random
from utils.config import get_config
from utils.evaluation import get_eval
from importlib import import_module
from torch.nn.modules.loss import CrossEntropyLoss
from monai.losses import DiceCELoss
from einops import rearrange
from models.model_dict import get_model
from utils.data_us import JointTransform2D, ImageToImage2D
from utils.loss_functions.sam_loss import get_criterion
from utils.generate_prompts import get_click_prompt
def main():
# ============================================================================= parameters setting ====================================================================================
parser = argparse.ArgumentParser(description='Networks')
parser.add_argument('--modelname', default='SAMUS', type=str, help='type of model, e.g., SAM, SAMFull, MedSAM, MSA, SAMed, SAMUS...')
parser.add_argument('-encoder_input_size', type=int, default=256, help='the image size of the encoder input, 1024 in SAM and MSA, 512 in SAMed, 256 in SAMUS')
parser.add_argument('-low_image_size', type=int, default=128, help='the image embedding size, 256 in SAM and MSA, 128 in SAMed and SAMUS')
parser.add_argument('--task', default='US30K', help='task or dataset name')
parser.add_argument('--vit_name', type=str, default='vit_b', help='select the vit model for the image encoder of sam')
parser.add_argument('--sam_ckpt', type=str, default='checkpoints/sam_vit_b_01ec64.pth', help='Pretrained checkpoint of SAM')
parser.add_argument('--batch_size', type=int, default=8, help='batch_size per gpu') # SAMed is 12 bs with 2n_gpu and lr is 0.005
parser.add_argument('--n_gpu', type=int, default=1, help='total gpu')
parser.add_argument('--base_lr', type=float, default=0.0005, help='segmentation network learning rate, 0.005 for SAMed, 0.0001 for MSA') #0.0006
parser.add_argument('--warmup', type=bool, default=False, help='If activated, warp up the learning from a lower lr to the base_lr')
parser.add_argument('--warmup_period', type=int, default=250, help='Warp up iterations, only valid whrn warmup is activated')
parser.add_argument('-keep_log', type=bool, default=False, help='keep the loss&lr&dice during training or not')
args = parser.parse_args()
opt = get_config(args.task)
device = torch.device(opt.device)
if args.keep_log:
logtimestr = time.strftime('%m%d%H%M') # initialize the tensorboard for record the training process
boardpath = opt.tensorboard_path + args.modelname + opt.save_path_code + logtimestr
if not os.path.isdir(boardpath):
os.makedirs(boardpath)
TensorWriter = SummaryWriter(boardpath)
# =============================================================== add the seed to make sure the results are reproducible ==============================================================
seed_value = 1234 # the number of seed
np.random.seed(seed_value) # set random seed for numpy
random.seed(seed_value) # set random seed for python
os.environ['PYTHONHASHSEED'] = str(seed_value) # avoid hash random
torch.manual_seed(seed_value) # set random seed for CPU
torch.cuda.manual_seed(seed_value) # set random seed for one GPU
torch.cuda.manual_seed_all(seed_value) # set random seed for all GPU
torch.backends.cudnn.deterministic = True # set random seed for convolution
# =========================================================================== model and data preparation ============================================================================
# register the sam model
model = get_model(args.modelname, args=args, opt=opt)
opt.batch_size = args.batch_size * args.n_gpu
tf_train = JointTransform2D(img_size=args.encoder_input_size, low_img_size=args.low_image_size, ori_size=opt.img_size, crop=opt.crop, p_flip=0.0, p_rota=0.5, p_scale=0.5, p_gaussn=0.0,
p_contr=0.5, p_gama=0.5, p_distor=0.0, color_jitter_params=None, long_mask=True) # image reprocessing
tf_val = JointTransform2D(img_size=args.encoder_input_size, low_img_size=args.low_image_size, ori_size=opt.img_size, crop=opt.crop, p_flip=0, color_jitter_params=None, long_mask=True)
train_dataset = ImageToImage2D(opt.data_path, opt.train_split, tf_train, img_size=args.encoder_input_size)
val_dataset = ImageToImage2D(opt.data_path, opt.val_split, tf_val, img_size=args.encoder_input_size) # return image, mask, and filename
trainloader = DataLoader(train_dataset, batch_size=opt.batch_size, shuffle=True, num_workers=8, pin_memory=True)
valloader = DataLoader(val_dataset, batch_size=opt.batch_size, shuffle=False, num_workers=8, pin_memory=True)
model.to(device)
if opt.pre_trained:
checkpoint = torch.load(opt.load_path)
new_state_dict = {}
for k,v in checkpoint.items():
if k[:7] == 'module.':
new_state_dict[k[7:]] = v
else:
new_state_dict[k] = v
model.load_state_dict(new_state_dict)
if args.n_gpu > 1:
model = nn.DataParallel(model)
if args.warmup:
b_lr = args.base_lr / args.warmup_period
optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=b_lr, betas=(0.9, 0.999), weight_decay=0.1)
else:
b_lr = args.base_lr
optimizer = optim.Adam(model.parameters(), lr=args.base_lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)
criterion = get_criterion(modelname=args.modelname, opt=opt)
pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print("Total_params: {}".format(pytorch_total_params))
# ========================================================================= begin to train the model ============================================================================
iter_num = 0
max_iterations = opt.epochs * len(trainloader)
best_dice, loss_log, dice_log = 0.0, np.zeros(opt.epochs+1), np.zeros(opt.epochs+1)
for epoch in range(opt.epochs):
# --------------------------------------------------------- training ---------------------------------------------------------
model.train()
train_losses = 0
for batch_idx, (datapack) in enumerate(trainloader):
imgs = datapack['image'].to(dtype = torch.float32, device=opt.device)
masks = datapack['low_mask'].to(dtype = torch.float32, device=opt.device)
bbox = torch.as_tensor(datapack['bbox'], dtype=torch.float32, device=opt.device)
pt = get_click_prompt(datapack, opt)
# -------------------------------------------------------- forward --------------------------------------------------------
pred = model(imgs, pt, bbox)
train_loss = criterion(pred, masks)
# -------------------------------------------------------- backward -------------------------------------------------------
optimizer.zero_grad()
train_loss.backward()
optimizer.step()
train_losses += train_loss.item()
print(train_loss)
# ------------------------------------------- adjust the learning rate when needed-----------------------------------------
if args.warmup and iter_num < args.warmup_period:
lr_ = args.base_lr * ((iter_num + 1) / args.warmup_period)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_
else:
if args.warmup:
shift_iter = iter_num - args.warmup_period
assert shift_iter >= 0, f'Shift iter is {shift_iter}, smaller than zero'
lr_ = args.base_lr * (1.0 - shift_iter / max_iterations) ** 0.9 # learning rate adjustment depends on the max iterations
for param_group in optimizer.param_groups:
param_group['lr'] = lr_
iter_num = iter_num + 1
# -------------------------------------------------- log the train progress --------------------------------------------------
print('epoch [{}/{}], train loss:{:.4f}'.format(epoch, opt.epochs, train_losses / (batch_idx + 1)))
if args.keep_log:
TensorWriter.add_scalar('train_loss', train_losses / (batch_idx + 1), epoch)
TensorWriter.add_scalar('learning rate', optimizer.state_dict()['param_groups'][0]['lr'], epoch)
loss_log[epoch] = train_losses / (batch_idx + 1)
# --------------------------------------------------------- evaluation ----------------------------------------------------------
if epoch % opt.eval_freq == 0:
model.eval()
dices, mean_dice, _, val_losses = get_eval(valloader, model, criterion=criterion, opt=opt, args=args)
print('epoch [{}/{}], val loss:{:.4f}'.format(epoch, opt.epochs, val_losses))
print('epoch [{}/{}], val dice:{:.4f}'.format(epoch, opt.epochs, mean_dice))
if args.keep_log:
TensorWriter.add_scalar('val_loss', val_losses, epoch)
TensorWriter.add_scalar('dices', mean_dice, epoch)
dice_log[epoch] = mean_dice
if mean_dice > best_dice:
best_dice = mean_dice
timestr = time.strftime('%m%d%H%M')
if not os.path.isdir(opt.save_path):
os.makedirs(opt.save_path)
save_path = opt.save_path + args.modelname + opt.save_path_code + '%s' % timestr + '_' + str(epoch) + '_' + str(best_dice)
torch.save(model.state_dict(), save_path + ".pth", _use_new_zipfile_serialization=False)
if epoch % opt.save_freq == 0 or epoch == (opt.epochs-1):
if not os.path.isdir(opt.save_path):
os.makedirs(opt.save_path)
save_path = opt.save_path + args.modelname + opt.save_path_code + '_' + str(epoch)
torch.save(model.state_dict(), save_path + ".pth", _use_new_zipfile_serialization=False)
if args.keep_log:
with open(opt.tensorboard_path + args.modelname + opt.save_path_code + logtimestr + '/trainloss.txt', 'w') as f:
for i in range(len(loss_log)):
f.write(str(loss_log[i])+'\n')
with open(opt.tensorboard_path + args.modelname + opt.save_path_code + logtimestr + '/dice.txt', 'w') as f:
for i in range(len(dice_log)):
f.write(str(dice_log[i])+'\n')
if __name__ == '__main__':
main()