Esempio n. 1
0
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)
Esempio n. 2
0
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)
Esempio n. 3
0
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
Esempio n. 4
0
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)
Esempio n. 5
0
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)
Esempio n. 6
0
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')