import os
import sys
from tqdm import tqdm
from tensorboardX import SummaryWriter
import argparse
import logging
import time
import random
import numpy as np
import torch
import torch.optim as optim
from torchvision import transforms
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
import torch.nn as nn
from torch.nn import BCEWithLogitsLoss, MSELoss
from torch.utils.data import DataLoader
from networks.vnet_sdf import VNet
from networks.discriminator import FC3DDiscriminator
from utils import ramps, losses
from dataloaders.la_heart import *
from dataloaders.utils import compute_sdf
parser = argparse.ArgumentParser()
parser.add_argument('--dataset_name', type=str, default='LA', help='dataset_name')
parser.add_argument('--root_path', type=str, default='/data/omnisky/postgraduate/Yb/data_set/LASet/data', help='Name of Experiment')
parser.add_argument('--exp', type=str, default='vnet', help='model_name')
parser.add_argument('--model', type=str, default='SASSNet', help='model_name')
parser.add_argument('--max_iterations', type=int, default=6000, help='maximum epoch number to train')
parser.add_argument('--batch_size', type=int, default=4, help='batch_size per gpu')
parser.add_argument('--labeled_bs', type=int, default=2, help='labeled_batch_size per gpu')
parser.add_argument('--base_lr', type=float, default=0.01, help='maximum epoch number to train')
parser.add_argument('--D_lr', type=float, default=1e-4, help='maximum discriminator learning rate to train')
parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training')
parser.add_argument('--labelnum', type=int, default=25, help='random seed') # 25有标签，98无标签
parser.add_argument('--max_samples', type=int, default=123, help='all samples')
parser.add_argument('--seed', type=int, default=1337, help='random seed')
parser.add_argument('--gpu', type=str, default='2', help='GPU to use')
parser.add_argument('--beta', type=float, default=0.3, help='balance factor to control regional and sdm loss')
parser.add_argument('--gamma', type=float, default=0.5, help='balance factor to control supervised and consistency loss')
### costs
parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay')
parser.add_argument('--consistency', type=float, default=0.01, help='consistency')
parser.add_argument('--consistency_rampup', type=float, default=40.0, help='consistency_rampup')
args = parser.parse_args()
num_classes = 2
patch_size = (112, 112, 80)
snapshot_path = "model/{}_{}_{}_labeled/{}".format(args.dataset_name, args.exp, args.labelnum, args.model)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
batch_size = args.batch_size * len(args.gpu.split(','))
max_iterations = args.max_iterations
base_lr = args.base_lr
labeled_bs = args.labeled_bs
if not args.deterministic:
cudnn.benchmark = True #
cudnn.deterministic = False #
else:
cudnn.benchmark = False # True #
cudnn.deterministic = True # False #
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
def cal_dice(output, target, eps=1e-3):
output = torch.sigmoid(output)
output = (output>0.5).float()
output = torch.squeeze(output)
inter = torch.sum(output * target) + eps
union = torch.sum(output) + torch.sum(target) + eps * 2
dice = 2 * inter / union
return dice
def get_current_consistency_weight(epoch):
# Consistency ramp-up from https://arxiv.org/abs/1610.02242
return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup)
if __name__ == "__main__":
# make logger file
if not os.path.exists(snapshot_path):
os.makedirs(snapshot_path)
logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO,
format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')
logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
logging.info(str(args))
model = VNet(n_channels=1, n_classes=num_classes-1, normalization='batchnorm', has_dropout=True).cuda()
D = FC3DDiscriminator(num_classes=num_classes - 1).cuda()
db_train = LAHeart(base_dir=args.root_path,
split='train',
transform=transforms.Compose([
RandomRotFlip(),
RandomCrop(patch_size),
ToTensor(),
]))
db_test = LAHeart(base_dir=args.root_path,
split='test',
transform=transforms.Compose([
CenterCrop(patch_size),
ToTensor()
]))
labeled_idxs = list(range(args.labelnum))
unlabeled_idxs = list(range(args.labelnum, args.max_samples))
batch_sampler = TwoStreamBatchSampler(labeled_idxs, unlabeled_idxs, batch_size, batch_size - labeled_bs)
def worker_init_fn(worker_id):
random.seed(args.seed+worker_id)
trainloader = DataLoader(db_train, batch_sampler=batch_sampler, num_workers=4, pin_memory=True,worker_init_fn=worker_init_fn)
test_loader = DataLoader(db_test, batch_size=1,shuffle=False, num_workers=4, pin_memory=True)
model.train()
Dopt = optim.Adam(D.parameters(), lr=args.D_lr, betas=(0.9,0.99))
optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=0.0001)
ce_loss = BCEWithLogitsLoss()
mse_loss = MSELoss()
writer = SummaryWriter(snapshot_path+'/log')
logging.info("{} itertations per epoch".format(len(trainloader)))
iter_num = 0
best_dice = 0
max_epoch = max_iterations//len(trainloader)+1
lr_ = base_lr
iterator = tqdm(range(max_epoch), ncols=70)
for epoch_num in iterator:
time1 = time.time()
for i_batch, sampled_batch in enumerate(trainloader):
time2 = time.time()
# print('fetch data cost {}'.format(time2-time1))
volume_batch, label_batch = sampled_batch['image'], sampled_batch['label']
volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda()
# Generate Discriminator target based on sampler
Dtarget = torch.tensor([1, 1, 0, 0]).cuda()
model.train()
D.eval()
outputs_tanh, outputs = model(volume_batch)
# print(outputs.shape)
outputs_soft = torch.sigmoid(outputs)
## calculate the loss
with torch.no_grad():
gt_dis = compute_sdf(label_batch[:].cpu().numpy(), outputs[:labeled_bs, 0, ...].shape)
gt_dis = torch.from_numpy(gt_dis).float().cuda()
loss_sdf = mse_loss(outputs_tanh[:labeled_bs, 0, ...], gt_dis)
loss_seg = ce_loss(outputs[:labeled_bs, 0, ...], label_batch[:labeled_bs].float())
loss_seg_dice = losses.dice_loss(outputs_soft[:labeled_bs, 0, :, :, :], label_batch[:labeled_bs] == 1)
supervised_loss = loss_seg_dice + args.beta * loss_sdf
consistency_weight = get_current_consistency_weight(iter_num//150)
Doutputs = D(outputs_tanh[labeled_bs:], volume_batch[labeled_bs:])
# G want D to misclassify unlabel data to label data.
loss_adv = F.cross_entropy(Doutputs, (Dtarget[:labeled_bs]).long())
loss = supervised_loss + consistency_weight*loss_adv
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Train D
model.eval()
D.train()
with torch.no_grad():
outputs_tanh, outputs = model(volume_batch)
Doutputs = D(outputs_tanh, volume_batch)
# D want to classify unlabel data and label data rightly.
D_loss = F.cross_entropy(Doutputs, Dtarget.long())
# Dtp and Dfn is unreliable because of the num of samples is small(4)
Dacc = torch.mean((torch.argmax(Doutputs, dim=1).float()==Dtarget.float()).float())
Dtp = torch.mean((torch.argmax(Doutputs, dim=1).float()==Dtarget.float()).float())
Dfn = torch.mean((torch.argmax(Doutputs, dim=1).float()==Dtarget.float()).float())
Dopt.zero_grad()
D_loss.backward()
Dopt.step()
iter_num = iter_num + 1
writer.add_scalar('lr', lr_, iter_num)
writer.add_scalar('loss/loss', loss, iter_num)
writer.add_scalar('loss/loss_seg', loss_seg, iter_num)
writer.add_scalar('loss/loss_dice', loss_seg_dice, iter_num)
writer.add_scalar('loss/loss_hausdorff', loss_sdf, iter_num)
writer.add_scalar('train/consistency_weight', consistency_weight, iter_num)
writer.add_scalar('loss/loss_adv', consistency_weight*loss_adv, iter_num)
writer.add_scalar('GAN/loss_adv', loss_adv, iter_num)
writer.add_scalar('GAN/D_loss', D_loss, iter_num)
writer.add_scalar('GAN/Dtp', Dtp, iter_num)
writer.add_scalar('GAN/Dfn', Dfn, iter_num)
logging.info(
'iteration %d : loss : %f, loss_weight: %f, loss_haus: %f, loss_seg: %f, loss_dice: %f' %
(iter_num, loss.item(), consistency_weight, loss_sdf.item(),
loss_seg.item(), loss_seg_dice.item()))
if iter_num >= 800 and iter_num % 200 == 0:
model.eval()
with torch.no_grad():
dice_sample = 0
for sampled_batch in test_loader:
img, lbl = sampled_batch['image'].cuda(), sampled_batch['label'].cuda()
_, outputs = model(img)
dice_once = cal_dice(outputs,lbl)
print(dice_once)
dice_sample += dice_once
dice_sample = dice_sample / len(test_loader)
print('Average center dice:{:.3f}'.format(dice_sample))
if dice_sample > best_dice:
best_dice = dice_sample
save_mode_path = os.path.join(snapshot_path, 'iter_{}_dice_{}.pth'.format(iter_num, best_dice))
save_best_path = os.path.join(snapshot_path, '{}_best_model.pth'.format(args.model))
torch.save(model.state_dict(), save_mode_path)
torch.save(model.state_dict(), save_best_path)
logging.info("save best model to {}".format(save_mode_path))
writer.add_scalar('Var_dice/Dice', dice_sample, iter_num)
writer.add_scalar('Var_dice/Best_dice', best_dice, iter_num)
model.train()
if iter_num >= max_iterations:
break
time1 = time.time()
if iter_num >= max_iterations:
iterator.close()
break
writer.close()
Datasets
Models
