--- a
+++ b/train_mean_teacher.py
@@ -0,0 +1,225 @@
+import os
+import sys
+from tqdm import tqdm
+from tensorboardX import SummaryWriter
+import argparse
+import logging
+import time
+import random
+
+import torch
+import torch.optim as optim
+from torchvision import transforms
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.utils.data import DataLoader
+
+from networks.vnet import VNet
+from utils import ramps, losses
+from dataloaders.la_heart import *
+
+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='MT', 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('--labelnum', type=int, default=25, help='trained samples')
+parser.add_argument('--max_samples', type=int, default=123, help='all samples')
+parser.add_argument('--base_lr', type=float, default=0.01, help='maximum epoch number to train')
+parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training')
+parser.add_argument('--seed', type=int, default=1337, help='random seed')
+parser.add_argument('--gpu', type=str, default='1', help='GPU to use')
+### costs
+parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay')
+parser.add_argument('--consistency_type', type=str, default="mse", help='consistency_type')
+parser.add_argument('--consistency', type=float, default=0.1, help='consistency')
+parser.add_argument('--consistency_rampup', type=float, default=40.0, help='consistency_rampup')
+args = parser.parse_args()
+
+
+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 args.deterministic:
+    cudnn.benchmark = False
+    cudnn.deterministic = True
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed(args.seed)
+
+num_classes = 2
+patch_size = (112, 112, 80)
+
+
+def cal_dice(output, target, eps=1e-3):
+    output = torch.argmax(output,dim=1)
+    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)
+
+
+def update_ema_variables(model, ema_model, alpha, global_step):
+    # Use the true average until the exponential average is more correct
+    alpha = min(1 - 1 / (global_step + 1), alpha)
+    for ema_param, param in zip(ema_model.parameters(), model.parameters()):
+        ema_param.data.mul_(alpha).add_(1 - alpha, param.data)
+
+
+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))
+
+
+    def create_model(ema=False):
+        # Network definition
+        net = VNet(n_channels=1, n_classes=num_classes, normalization='batchnorm', has_dropout=True)
+        model = net.cuda()
+        if ema:
+            for param in model.parameters():
+                param.detach_()
+        return model
+
+
+    model = create_model()
+    ema_model = create_model(ema=True)
+    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)
+
+
+    train_loader = 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()
+    ema_model.train()
+    optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=0.0001)
+
+    if args.consistency_type == 'mse':
+        consistency_criterion = losses.softmax_mse_loss
+    elif args.consistency_type == 'kl':
+        consistency_criterion = losses.softmax_kl_loss
+    else:
+        assert False, args.consistency_type
+
+    writer = SummaryWriter(snapshot_path + '/log')
+    logging.info("{} itertations per epoch".format(len(train_loader)))
+
+    iter_num = 0
+    best_dice = 0
+    max_epoch = max_iterations // len(train_loader) + 1
+    lr_ = base_lr
+    model.train()
+    for epoch_num in tqdm(range(max_epoch), ncols=70):
+        time1 = time.time()
+        for i_batch, sampled_batch in enumerate(train_loader):
+            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()
+            unlabeled_volume_batch = volume_batch[labeled_bs:]
+
+            noise = torch.clamp(torch.randn_like(unlabeled_volume_batch) * 0.1, -0.2, 0.2)
+            ema_inputs = unlabeled_volume_batch + noise
+            outputs = model(volume_batch)
+            with torch.no_grad():
+                ema_output = ema_model(ema_inputs)
+
+            # calculate the loss
+            loss_seg = F.cross_entropy(outputs[:labeled_bs], label_batch[:labeled_bs])
+            outputs_soft = F.softmax(outputs, dim=1)
+            loss_seg_dice = losses.dice_loss(outputs_soft[:labeled_bs, 1, :, :, :], label_batch[:labeled_bs] == 1)
+            supervised_loss = 0.5 * (loss_seg + loss_seg_dice)
+
+            consistency_weight = get_current_consistency_weight(iter_num // 150)
+            consistency_dist = consistency_criterion(outputs[labeled_bs:], ema_output) # (batch, 2, 112,112,80)
+            consistency_loss = consistency_weight * consistency_dist
+            loss = supervised_loss + consistency_loss
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            update_ema_variables(model, ema_model, args.ema_decay, iter_num)
+
+            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_seg_dice', loss_seg_dice, iter_num)
+            writer.add_scalar('train/consistency_weight', consistency_weight, iter_num)
+            writer.add_scalar('train/consistency_dist', consistency_dist, iter_num)
+
+            logging.info('iteration %d : loss : %f cons_dist: %f, loss_weight: %f' %
+                         (iter_num, loss.item(), consistency_dist.item(), consistency_weight))
+
+            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)
+                        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:
+            break
+    save_mode_path = os.path.join(snapshot_path, 'iter_' + str(max_iterations) + '.pth')
+    torch.save(model.state_dict(), save_mode_path)
+    logging.info("save model to {}".format(save_mode_path))
+    writer.close()