(args, model, snapshot_path)
| 38 | return performance |
| 39 | |
| 40 | def trainer_synapse(args, model, snapshot_path): |
| 41 | logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, |
| 42 | format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') |
| 43 | logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) |
| 44 | logging.info(str(args)) |
| 45 | base_lr = args.base_lr |
| 46 | num_classes = args.num_classes |
| 47 | batch_size = args.batch_size * args.n_gpu |
| 48 | |
| 49 | db_train = Synapse_dataset(base_dir=args.root_path, list_dir=args.list_dir, split="train", nclass=args.num_classes, |
| 50 | transform=transforms.Compose( |
| 51 | [RandomGenerator(output_size=[args.img_size, args.img_size])])) |
| 52 | |
| 53 | print("The length of train set is: {}".format(len(db_train))) |
| 54 | |
| 55 | def worker_init_fn(worker_id): |
| 56 | random.seed(args.seed + worker_id) |
| 57 | |
| 58 | trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True, |
| 59 | worker_init_fn=worker_init_fn) |
| 60 | |
| 61 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 62 | if torch.cuda.device_count() > 1 and args.n_gpu > 1: |
| 63 | print("Let's use", torch.cuda.device_count(), "GPUs!") |
| 64 | model = nn.DataParallel(model) |
| 65 | model.to(device) |
| 66 | |
| 67 | model.train() |
| 68 | ce_loss = CrossEntropyLoss() |
| 69 | dice_loss = DiceLoss(num_classes) |
| 70 | |
| 71 | #optimizer = optim.SGD(model.parameters(), lr=base_lr, momentum=0.9, weight_decay=0.0001) |
| 72 | optimizer = optim.AdamW(model.parameters(), lr=base_lr, weight_decay=0.0001) |
| 73 | writer = SummaryWriter(snapshot_path + '/log') |
| 74 | iter_num = 0 |
| 75 | max_epoch = args.max_epochs |
| 76 | max_iterations = args.max_epochs * len(trainloader) |
| 77 | logging.info("{} iterations per epoch. {} max iterations ".format(len(trainloader), max_iterations)) |
| 78 | best_performance = 0.0 |
| 79 | iterator = tqdm(range(max_epoch), ncols=70) |
| 80 | |
| 81 | for epoch_num in iterator: |
| 82 | |
| 83 | for i_batch, sampled_batch in enumerate(trainloader): |
| 84 | image_batch, label_batch = sampled_batch['image'], sampled_batch['label'] |
| 85 | image_batch, label_batch = image_batch.cuda(), label_batch.squeeze(1).cuda() |
| 86 | |
| 87 | P = model(image_batch, mode='train') |
| 88 | |
| 89 | if not isinstance(P, list): |
| 90 | P = [P] |
| 91 | if epoch_num == 0 and i_batch == 0: |
| 92 | n_outs = len(P) |
| 93 | out_idxs = list(np.arange(n_outs)) #[0, 1, 2, 3]#, 4, 5, 6, 7] |
| 94 | if args.supervision == 'mutation': |
| 95 | ss = [x for x in powerset(out_idxs)] |
| 96 | elif args.supervision == 'deep_supervision': |
| 97 | ss = [[x] for x in out_idxs] |
nothing calls this directly
no test coverage detected