def train(train_loader, model, optimizer, epoch):
model.train()
# ---- multi-scale training ----
loss_record = AvgMeter()
for i, pack in enumerate(train_loader, start=1):
optimizer.zero_grad()
# ---- data prepare ----
images, gts = pack['image'], pack['label']
images = Variable(images).cuda()
gts = Variable(gts).cuda()
# ---- forward ----
pred = model(images)
# ---- loss function ----
loss4 = BCEDiceLoss(pred[4], gts)
loss3 = BCEDiceLoss(pred[3], gts)
loss2 = BCEDiceLoss(pred[2], gts)
loss1 = BCEDiceLoss(pred[1], gts)
loss0 = BCEDiceLoss(pred[0], gts)
loss = loss0 + loss1 + loss2 + loss3 + loss4
# ---- backward ----
loss.backward()
optimizer.step()
# ---- recording loss ----
loss_record.update(loss.data, opt.batchsize)
# ---- train visualization ----
if i == total_step:
print('{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], '
'[loss: {:.4f}]'.format(datetime.now(), epoch, opt.epoch, i,
total_step, loss_record.show()))
save_path = './snapshots/{}/'.format(opt.save_root)
os.makedirs(save_path, exist_ok=True)
if epoch % 50 == 0:
torch.save(model.state_dict(), save_path + 'model-%d.pth' % epoch)
print('[Saving Snapshot:]', save_path + 'model-%d.pth' % epoch)
def train(train_loader, model, optimizer, epochs, batch_size, train_size, clip,
test_path):
best_dice_score = 0
for epoch in range(1, epochs): # 99 epoch
adjust_lr(optimizer, lr, epoch, 0.1, 200)
model.train()
size_rates = [0.75, 1, 1.25]
loss_record = AvgMeter()
criterion = WIoUBCELoss()
for i, pack in enumerate(train_loader, start=1):
for rate in size_rates:
optimizer.zero_grad()
images, gts = pack
images = Variable(images).cuda()
gts = Variable(gts).cuda()
trainsize = int(round(train_size * rate / 32) * 32)
if rate != 1:
images = F.upsample(images,
size=(trainsize, trainsize),
mode='bilinear',
align_corners=True)
gts = F.upsample(gts,
size=(trainsize, trainsize),
mode='bilinear',
align_corners=True)
# predict
predict_maps = model(images)
loss = criterion(predict_maps, gts)
loss.backward()
clip_gradient(optimizer, clip)
optimizer.step()
if rate == 1:
loss_record.update(loss.data, batch_size)
if i % 20 == 0 or i == total_step:
print(
f'{datetime.now()} Epoch [{epoch}/{epochs}], Step [{i}/{total_step}], Loss: {loss_record.show()}'
)
train_logger.info(
f'{datetime.now()} Epoch [{epoch}/{epochs}], Step [{i}/{total_step}], Loss: {loss_record.show()}'
)
save_path = 'checkpoints/'
os.makedirs(save_path, exist_ok=True)
if (epoch + 1) % 1 == 0:
meandice = validation(model, test_path)
print(f'meandice: {meandice}')
train_logger.info(f'meandice: {meandice}')
if meandice > best_dice_score:
best_dice_score = meandice
torch.save(model.state_dict(), save_path + 'effnetv2pd.pth')
print('[Saving Snapshots:]', save_path + 'effnetv2pd.pth',
meandice)
if epoch in [50, 60, 70]:
file_ = 'effnetv2pd_' + str(epoch) + '.pth'
torch.save(model.state_dict(), save_path + file_)
print('[Saving Snapshots:]', save_path + file_, meandice)
def train(train_loader, model, optimizer, epoch, test_path, best):
model.train()
size_rates = [0.75, 1, 1.25]
loss_attention_record, loss_detection_record = AvgMeter(), AvgMeter()
criterion = WIoUBCELoss()
for i, pack in enumerate(train_loader, start=1):
for rate in size_rates:
optimizer.zero_grad()
images, gts = pack
images = Variable(images).cuda()
gts = Variable(gts).cuda()
trainsize = int(round(opt.trainsize * rate / 32) * 32)
if rate != 1:
images = F.upsample(images,
size=(trainsize, trainsize),
mode='bilinear',
align_corners=True)
gts = F.upsample(gts,
size=(trainsize, trainsize),
mode='bilinear',
align_corners=True)
attention_map, detection_map = model(images)
loss1 = criterion(attention_map, gts)
loss2 = criterion(detection_map, gts)
loss = loss1 + loss2
loss.backward()
clip_gradient(optimizer, opt.clip)
optimizer.step()
if rate == 1:
loss_attention_record.update(loss1.data, opt.batchsize)
loss_detection_record.update(loss2.data, opt.batchsize)
if i % 20 == 0 or i == total_step:
print(
f'{datetime.now()} Epoch [{epoch}/{opt.epoch}], Step [{i}/{total_step}], \
Loss [attention_loss: {loss_attention_record.show()}, detection_loss: {loss_detection_record.show()}]'
)
train_logger.info(
f'{datetime.now()} Epoch [{epoch}/{opt.epoch}], Step [{i}/{total_step}], \
Loss [attention_loss: {loss_attention_record.show()}, detection_loss: {loss_detection_record.show()}]'
)
save_path = 'snapshots/{}/'.format(opt.train_save)
os.makedirs(save_path, exist_ok=True)
if (epoch + 1) % 1 == 0:
meandice = validation(model, test_path)
if meandice > best:
best = meandice
torch.save(model.state_dict(), save_path + 'HarDMSEG-best.pth')
print('[Saving Snapshots:]', save_path + 'HarDMSEG-best.pth',
meandice)
train_logger.info(
f'[Saving Snapshots: {save_path + "HarGMSEG-best"} {meandice}]'
)
return best
def train(train_loader, model, optimizer, epochs, batch_size, train_size, clip, test_path):
best_dice_score = 0
for epoch in range(1, epochs):
adjust_lr(optimizer, lr, epoch, 0.1, 200)
for param in optimizer.param_groups:
print(param['lr'])
model.train()
size_rates = [0.75, 1, 1.25]
loss_record2, loss_record3, loss_record4, loss_record5 = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter()
criterion = WIoUBCELoss()
for i, pack in enumerate(train_loader, start=1):
for rate in size_rates:
optimizer.zero_grad()
images, gts = pack
images = Variable(images).cuda()
gts = Variable(gts).cuda()
trainsize = int(round(train_size * rate / 32) * 32)
if rate != 1:
images = F.upsample(images, size=(
trainsize, trainsize), mode='bilinear', align_corners=True)
gts = F.upsample(gts, size=(
trainsize, trainsize), mode='bilinear', align_corners=True)
# predict
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2 = model(
images)
loss5 = criterion(lateral_map_5, gts)
loss4 = criterion(lateral_map_4, gts)
loss3 = criterion(lateral_map_3, gts)
loss2 = criterion(lateral_map_2, gts)
loss = loss2 + loss3 + loss4 + loss5
loss.backward()
clip_gradient(optimizer, clip)
optimizer.step()
if rate == 1:
loss_record2.update(loss2.data, batch_size)
loss_record3.update(loss3.data, batch_size)
loss_record4.update(loss4.data, batch_size)
loss_record5.update(loss5.data, batch_size)
if i % 20 == 0 or i == total_step:
print(f'{datetime.now()} Epoch [{epoch}/{epochs}], Step [{i}/{total_step}], [lateral-2: {loss_record2.show()}, lateral-3: {loss_record3.show()}, lateral-4: {loss_record4.show()}, lateral-5: {loss_record5.show()},]')
train_logger.info(
f'{datetime.now()} Epoch [{epoch}/{epochs}], Step [{i}/{total_step}], [lateral-2: {loss_record2.show()}, lateral-3: {loss_record3.show()}, lateral-4: {loss_record4.show()}, lateral-5: {loss_record5.show()},]')
save_path = 'checkpoints/'
os.makedirs(save_path, exist_ok=True)
if (epoch+1) % 1 == 0:
meandice = validation(model, test_path)
print(f'meandice: {meandice}')
train_logger.info(f'meandice: {meandice}')
if meandice > best_dice_score:
best_dice_score = meandice
torch.save(model.state_dict(), save_path + 'PraHarDNet.pth')
print('[Saving Snapshots:]', save_path + 'PraHarDNet.pth', meandice)
def train(train_loader, model, optimizer, epoch):
model.train()
# ---- multi-scale training ----
size_rates = [0.75, 1, 1.25]
loss_record2, loss_record3, loss_record4, loss_record5 = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter()
for i, pack in enumerate(train_loader, start=1):
for rate in size_rates:
optimizer.zero_grad()
# ---- data prepare ----
images, gts = pack
images = Variable(images).cuda()
gts = Variable(gts).cuda()
# ---- rescale ----
trainsize = int(round(opt.trainsize * rate / 32) * 32)
if rate != 1:
images = F.upsample(images,
size=(trainsize, trainsize),
mode='bilinear',
align_corners=True)
gts = F.upsample(gts,
size=(trainsize, trainsize),
mode='bilinear',
align_corners=True)
# ---- forward ----
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2 = model(
images)
# ---- loss function ----
loss5 = structure_loss(lateral_map_5, gts)
loss4 = structure_loss(lateral_map_4, gts)
loss3 = structure_loss(lateral_map_3, gts)
loss2 = structure_loss(lateral_map_2, gts)
loss = loss2 + loss3 + loss4 + loss5 # TODO: try different weights for loss
# ---- backward ----
loss.backward()
clip_gradient(optimizer, opt.clip)
optimizer.step()
# ---- recording loss ----
if rate == 1:
loss_record2.update(loss2.data, opt.batchsize)
loss_record3.update(loss3.data, opt.batchsize)
loss_record4.update(loss4.data, opt.batchsize)
loss_record5.update(loss5.data, opt.batchsize)
# ---- train visualization ----
if i % 20 == 0 or i == total_step:
print(
'{} Epoch [{:03d}/{:03d}], Step [{:04d}/{:04d}], '
'[lateral-2: {:.4f}, lateral-3: {:0.4f}, lateral-4: {:0.4f}, lateral-5: {:0.4f}]'
.format(datetime.now(), epoch, opt.epoch, i, total_step,
loss_record2.show(), loss_record3.show(),
loss_record4.show(), loss_record5.show()))
save_path = 'snapshots/{}/'.format(opt.train_save)
os.makedirs(save_path, exist_ok=True)
if (epoch + 1) % 10 == 0:
torch.save(model.state_dict(), save_path + 'PraNet-%d.pth' % epoch)
print('[Saving Snapshot:]', save_path + 'PraNet-%d.pth' % epoch)
def main():
parser = ArgumentParser()
parser.add_argument("-c",
"--config",
required=True,
default="configs/default_config.yaml")
args = parser.parse_args()
logger.info("Loading config")
config_path = args.config
config = load_cfg(config_path)
gts = []
prs = []
folds = config["test"]["folds"]
print(folds)
dataset = config["dataset"]["test_data_path"][0].split("/")[-1]
if len(folds.keys()) == 1:
logger.add(
f'logs/test_{config["model"]["arch"]}_{str(datetime.now())}_{list(folds.keys())[0]}_{dataset}.log',
rotation="10 MB",
)
else:
logger.add(
f'logs/test_{config["model"]["arch"]}_{str(datetime.now())}_kfold.log',
rotation="10 MB",
)
for id in list(folds.keys()):
test_img_paths = []
test_mask_paths = []
test_data_path = config["dataset"]["test_data_path"]
for i in test_data_path:
test_img_paths.extend(glob(os.path.join(i, "*")))
test_mask_paths.extend(glob(os.path.join(i, "*")))
test_img_paths.sort()
test_mask_paths.sort()
test_transform = None
test_loader = get_loader(
test_img_paths,
test_mask_paths,
transform=test_transform,
**config["test"]["dataloader"],
type="test",
)
test_size = len(test_loader)
epochs = folds[id]
if type(epochs) != list:
epochs = [3 * (epochs // 3) + 2]
elif len(epochs) == 2:
epochs = [i for i in range(epochs[0], epochs[1])]
# epochs = [3 * i + 2 for i in range(epochs[0] // 3, (epochs[1] + 1) // 3)]
elif len(epochs) == 1:
epochs = [3 * (epochs[0] // 3) + 2]
else:
logger.debug("Model path must have 0 or 1 num")
break
for e in epochs:
# MODEL
logger.info("Loading model")
model_prams = config["model"]
import network.models as models
arch = model_prams["arch"]
model = models.__dict__[arch]() # Pranet
if "save_dir" not in model_prams:
save_dir = os.path.join("snapshots",
model_prams["arch"] + "_kfold")
else:
save_dir = config["model"]["save_dir"]
model_path = os.path.join(
save_dir,
f"PraNetDG-fold{id}-{e}.pth",
)
device = torch.device("cpu")
# model.cpu()
model.cuda()
model.eval()
logger.info(f"Loading from {model_path}")
try:
model.load_state_dict(
torch.load(model_path)["model_state_dict"])
except RuntimeError:
model.load_state_dict(torch.load(model_path))
test_fold = "fold" + str(config["dataset"]["fold"])
logger.info(f"Start testing fold{id} epoch {e}")
if "visualize_dir" not in config["test"]:
visualize_dir = "results"
else:
visualize_dir = config["test"]["visualize_dir"]
test_fold = "fold" + str(id)
logger.info(
f"Start testing {len(test_loader)} images in {dataset} dataset"
)
vals = AvgMeter()
H, W, T = 240, 240, 155
for i, pack in tqdm.tqdm(enumerate(test_loader, start=1)):
image, gt, filename, img = pack
name = os.path.splitext(filename[0])[0]
ext = os.path.splitext(filename[0])[1]
# print(gt.shape,image.shape,"ppp")
# import sys
# sys.exit()
gt = gt[0]
gt = np.asarray(gt, np.float32)
res2 = 0
image = image.cuda()
res5, res4, res3, res2 = model(image)
# res = res2
# res = F.upsample(
# res, size=gt.shape, mode="bilinear", align_corners=False
# )
# res = res.sigmoid().data.cpu().numpy().squeeze()
# res = (res - res.min()) / (res.max() - res.min() + 1e-8)
output = res2[0, :, :H, :W, :T].cpu().detach().numpy()
output = output.argmax(
0
) # (num_classes,height,width,depth) num_classes is now one-hot
target_cpu = gt[:H, :W, :T].numpy()
scores = softmax_output_dice(output, target_cpu)
vals.update(np.array(scores))
# msg += ', '.join(['{}: {:.4f}'.format(k, v) for k, v in zip(keys, scores)])
seg_img = np.zeros(shape=(H, W, T), dtype=np.uint8)
# same as res.round()
seg_img[np.where(output == 1)] = 1
seg_img[np.where(output == 2)] = 2
seg_img[np.where(output == 3)] = 4
# if verbose:
logger.info(
f'1:{np.sum(seg_img==1)} | 2: {np.sum(seg_img==2)} | 4: {np.sum(seg_img==4)}'
)
logger.info(
f'WT: {np.sum((seg_img==1)|(seg_img==2)|(seg_img==4))} | TC: {np.sum((seg_img==1)|(seg_img==4))} | ET: {np.sum(seg_img==4)}'
)
overwrite = config["test"]["vis_overwrite"]
vis_x = config["test"]["vis_x"]
if config["test"]["visualize"]:
oname = os.path.join(visualize_dir, 'submission',
name[:-8] + '_pred.nii.gz')
save_img(
oname,
seg_img,
"nib",
overwrite,
)
logger.info(vals.avg)
def fit(self,
train_loader,
is_val=False,
test_loader=None,
img_size=352,
start_from=0,
num_epochs=200,
batchsize=16,
clip=0.5,
fold=4):
size_rates = [0.75, 1, 1.25]
test_fold = f'fold{fold}'
start = timeit.default_timer()
for epoch in range(start_from, num_epochs):
self.net.train()
loss_all, loss_record2, loss_record3, loss_record4, loss_record5 = AvgMeter(
), AvgMeter(), AvgMeter(), AvgMeter(), AvgMeter()
for i, pack in enumerate(train_loader, start=1):
for rate in size_rates:
self.optimizer.zero_grad()
# ---- data prepare ----
images, gts = pack
# images, gts, paths, oriimgs = pack
images = Variable(images).cuda()
gts = Variable(gts).cuda()
trainsize = int(round(img_size * rate / 32) * 32)
if rate != 1:
images = F.upsample(images,
size=(trainsize, trainsize),
mode='bilinear',
align_corners=True)
gts = F.upsample(gts,
size=(trainsize, trainsize),
mode='bilinear',
align_corners=True)
lateral_map_5, lateral_map_4, lateral_map_3, lateral_map_2 = self.net(
images)
# lateral_map_5 = self.net(images)
loss5 = self.loss(lateral_map_5, gts)
# loss4 = self.loss(lateral_map_4, gts)
# loss3 = self.loss(lateral_map_3, gts)
# loss2 = self.loss(lateral_map_2, gts)
# loss = loss2 + loss3 + loss4 + loss5
loss = loss5
loss.backward()
clip_gradient(self.optimizer, clip)
self.optimizer.step()
if rate == 1:
# loss_record2.update(loss2.data, batchsize)
# loss_record3.update(loss3.data, batchsize)
# loss_record4.update(loss4.data, batchsize)
loss_record5.update(loss5.data, batchsize)
loss_all.update(loss.data, batchsize)
# self.writer.add_scalar("Loss2", loss_record2.show(), (epoch-1)*len(train_loader) + i)
# self.writer.add_scalar("Loss3", loss_record3.show(), (epoch-1)*len(train_loader) + i)
# self.writer.add_scalar("Loss4", loss_record4.show(), (epoch-1)*len(train_loader) + i)
self.writer.add_scalar(
"Loss5", loss_record5.show(),
(epoch - 1) * len(train_loader) + i)
self.writer.add_scalar(
"Loss", loss_all.show(),
(epoch - 1) * len(train_loader) + i)
total_step = len(train_loader)
if i % 25 == 0 or i == total_step:
# self.logger.info('{} Epoch [{:03d}/{:03d}], with lr = {}, Step [{:04d}/{:04d}],\
# [loss_record2: {:.4f},loss_record3: {:.4f},loss_record4: {:.4f},loss_record5: {:.4f}]'.
# format(datetime.now(), epoch, epoch, self.optimizer.param_groups[0]["lr"],i, total_step,\
# loss_record2.show(), loss_record3.show(), loss_record4.show(), loss_record5.show()
# ))
self.logger.info(
'{} Epoch [{:03d}/{:03d}], with lr = {}, Step [{:04d}/{:04d}],\
[loss_record5: {:.4f}]'.format(
datetime.now(), epoch, epoch,
self.optimizer.param_groups[0]["lr"], i,
total_step, loss_record5.show()))
if (is_val):
self.val(test_loader, epoch)
os.makedirs(self.save_dir, exist_ok=True)
if (epoch + 1) % 3 == 0 and epoch > self.save_from or epoch == 23:
torch.save(
{
"model_state_dict": self.net.state_dict(),
"lr": self.optimizer.param_groups[0]["lr"]
},
os.path.join(self.save_dir,
'PraNetDG-' + test_fold + '-%d.pth' % epoch))
self.logger.info(
'[Saving Snapshot:]' +
os.path.join(self.save_dir, 'PraNetDG-' + test_fold +
'-%d.pth' % epoch))
self.scheduler.step()
self.writer.flush()
self.writer.close()
end = timeit.default_timer()
self.logger.info("Training cost: " + str(end - start) + 'seconds')