#########
                 #inference
                 if crop_img.dim == 3:
                     raw_crop_ = raw_crop[np.newaxis, :, :, :]
                 else:
                     raw_crop_ = raw_crop[np.newaxis, np.newaxis, :, :]
                 inputs = torch.Tensor(raw_crop_).to(device)
                 inputs = F.pad(inputs, (PAD, PAD, PAD, PAD))
                 with torch.no_grad():
                     pred = model(inputs)
                 # pred = pred[0]
                 pred = F.pad(pred, (-PAD, -PAD, -PAD, -PAD))
                 # pred = F.softmax(pred, dim=1)
                 # pred = torch.argmax(pred, dim=1).squeeze(0)
                 pred = pred.data.cpu().numpy()
                 pred = np.squeeze(pred)
                 #########
                 crop_img.save(i, j, pred)
         results = crop_img.result()
         if thresd is None:
             results[results<=0] = 0
             results[results>1] = 1
         else:
             results[results<=thresd] = 0
             results[results>thresd] = 1
         results = (results * 255).astype(np.uint8)
         # if cfg.TRAIN.track == 'complex':
         #     results = results[141:141+9959, 141:141+9958]
         print('COST TIME: ', (time.time() - start))
         cv2.imwrite(os.path.join(save_path, f_img[:-3]+'tiff'), results)
 print('***Done***')
def loop(cfg, train_provider, valid_provider, model, criterion, optimizer, iters, writer):
    PAD = cfg.TRAIN.pad
    f_loss_txt = open(os.path.join(cfg.record_path, 'loss.txt'), 'a')
    f_valid_txt = open(os.path.join(cfg.record_path, 'valid.txt'), 'a')
    rcd_time = []
    sum_time = 0
    sum_loss = 0
    valid_score = []
    valid_score_tmp = []
    thresd = cfg.TRAIN.thresd
    most_f1 = 0
    most_iters = 0
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    while iters <= cfg.TRAIN.total_iters:
        # train
        iters += 1
        t1 = time.time()
        input, target = train_provider.next()
        
        # decay learning rate
        if cfg.TRAIN.end_lr == cfg.TRAIN.base_lr:
            current_lr = cfg.TRAIN.base_lr
        else:
            current_lr = calculate_lr(iters)
            for param_group in optimizer.param_groups:
                param_group['lr'] = current_lr
        
        optimizer.zero_grad()
        input = F.pad(input, (PAD, PAD, PAD, PAD))
        pred = model(input)
        pred = F.pad(pred, (-PAD, -PAD, -PAD, -PAD))
        # target = torch.unsqueeze(target, dim=1)
        # if iters == 1:
        #     writer.add_graph(model, (input,))

        loss = criterion(pred, target)
        loss.backward()
        if cfg.TRAIN.weight_decay is not None:
            for group in optimizer.param_groups:
                for param in group['params']:
                    param.data = param.data.add(-cfg.TRAIN.weight_decay * group['lr'], param.data)
        optimizer.step()
        
        sum_loss += loss.item()
        sum_time += time.time() - t1
        
        # log train
        if iters % cfg.TRAIN.display_freq == 0:
            rcd_time.append(sum_time)
            logging.info('step %d, loss = %.4f (wt: *10, lr: %.8f, et: %.2f sec, rd: %.2f min)'
                            % (iters, sum_loss / cfg.TRAIN.display_freq * 10, current_lr, sum_time,
                            (cfg.TRAIN.total_iters - iters) / cfg.TRAIN.display_freq * np.mean(np.asarray(rcd_time)) / 60))
            writer.add_scalar('loss', sum_loss / cfg.TRAIN.display_freq * 10, iters)
            f_loss_txt.write('step = ' + str(iters) + ', loss = ' + str(sum_loss / cfg.TRAIN.display_freq * 10))
            f_loss_txt.write('\n')
            f_loss_txt.flush()
            sys.stdout.flush()
            sum_time = 0
            sum_loss = 0
        
        # valid
        if iters % cfg.TRAIN.valid_freq == 0:
            input = F.pad(input, (-PAD, -PAD, -PAD, -PAD))
            input0 = (np.squeeze(input[0].data.cpu().numpy()) * 255).astype(np.uint8)
            target = (np.squeeze(target[0].data.cpu().numpy()) * 255).astype(np.uint8)
            pred = np.squeeze(pred[0].data.cpu().numpy())
            pred[pred>1] = 1; pred[pred<0] = 0
            pred = (pred * 255).astype(np.uint8)
            input0 = input0[0]
            im_cat = np.concatenate([input0, pred, target], axis=1)
            Image.fromarray(im_cat).save(os.path.join(cfg.cache_path, '%06d.png' % iters))
        
        # save
        if iters % cfg.TRAIN.save_freq == 0:
            if cfg.TRAIN.if_valid:
                for k in range(valid_provider.data.num):
                    raw, label = valid_provider.data.gen(k)
                    crop_img = Crop_image(raw,crop_size=cfg.TRAIN.crop_size,overlap=cfg.TRAIN.overlap)
                    for i in range(crop_img.num):
                        for j in range(crop_img.num):
                            raw_crop = crop_img.gen(i, j)
                            #########
                            #inference
                            if crop_img.dim == 3:
                                raw_crop_ = raw_crop[np.newaxis, :, :, :]
                            else:
                                raw_crop_ = raw_crop[np.newaxis, np.newaxis, :, :]
                            inputs = torch.Tensor(raw_crop_).to(device)
                            inputs = F.pad(inputs, (PAD, PAD, PAD, PAD))
                            with torch.no_grad():
                                pred = model(inputs)
                            pred = F.pad(pred, (-PAD, -PAD, -PAD, -PAD))
                            pred = pred.data.cpu().numpy()
                            pred = np.squeeze(pred)
                            #########
                            crop_img.save(i, j, pred)
                    results = crop_img.result()
                    results[results<=thresd] = 0
                    results[results>thresd] = 1
                    temp_label = label.flatten()
                    temp_result = results.flatten()
                    f1_common = f1_score(1 - temp_label, 1 - temp_result)
                    f1 = 0
                    results_img = (results * 255).astype(np.uint8)
                    label_img = (label * 255).astype(np.uint8)
                    im_cat_valid = np.concatenate([results_img, label_img], axis=1)
                    if k == valid_provider.data.num - 1:
                        Image.fromarray(im_cat_valid).save(os.path.join(cfg.cache_path, 'valid_%06d.png' % iters))
                    valid_score.append(f1)
                    valid_score_tmp.append(f1_common)
                avg_f1_soft = sum(valid_score) / len(valid_score)
                avg_f1_tmp = sum(valid_score_tmp) / len(valid_score_tmp)
                logging.info('step %d, f1_soft = %.6f' % (iters, avg_f1_soft))
                logging.info('step %d, f1_common = %.6f' % (iters, avg_f1_tmp))
                writer.add_scalar('valid', avg_f1_tmp, iters)
                f_valid_txt.write('step %d, f1 = %.6f' % (iters, avg_f1_tmp))
                f_valid_txt.write('\n')
                f_valid_txt.flush()
                sys.stdout.flush()
                if avg_f1_tmp > most_f1:
                    most_f1 = avg_f1_tmp
                    most_iters = iters
                    states = {'current_iter': iters, 'valid_result': avg_f1_tmp,
                            'model_weights': model.state_dict()}
                    torch.save(states, os.path.join(cfg.save_path, 'model.ckpt'))
                    print('***************save modol, when f1 = %.6f and iters = %d.***************' % (most_f1, most_iters))
            else:
                states = {'current_iter': iters, 'valid_result': None,
                        'model_weights': model.state_dict()}
                torch.save(states, os.path.join(cfg.save_path, 'model-%d.ckpt' % iters))
                print('save model-%d' % iters)
Example #3
0
def loop(cfg, train_provider, valid_provider, model, criterion, optimizer, iters, writer):
    PAD = cfg.TRAIN.pad
    f_loss_txt = open(os.path.join(cfg.record_path, 'loss.txt'), 'w')
    f_valid_txt = open(os.path.join(cfg.record_path, 'valid.txt'), 'w')
    rcd_time = []
    sum_time = 0
    sum_loss = 0
    valid_score = []
    valid_score_tmp = []
    thresd = cfg.TRAIN.thresd
    most_f1 = 0
    most_iters = 0
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    ######################################
    # loss_vgg = PerceptualLoss(mode=2).to(device)
    model_vgg = VGG19().to(device)
    cuda_count = torch.cuda.device_count()
    if cuda_count > 1:
        model_vgg = nn.DataParallel(model_vgg)
        print('VGG build on %d GPUs ... ' % cuda_count, flush=True)
    else:
        print('VGG build on a single GPU ... ', flush=True)
    # cuda_count = torch.cuda.device_count()
    # if cuda_count > 1:
    #     if cfg.TRAIN.batch_size % cuda_count == 0:
    #         loss_vgg = nn.DataParallel(loss_vgg)
    ######################################
    # mse_weight = cfg.TRAIN.mse_weight
    vgg_weight = cfg.TRAIN.vgg_weight

    ### loss function
    # if cfg.MODEL.loss_func_logits:
    #     loss_function = F.binary_cross_entropy_with_logits
    # else:
    #     loss_function = F.binary_cross_entropy

    # final_loss = 0
    while iters <= cfg.TRAIN.total_iters:
        
        # train
        iters += 1
        t1 = time.time()
        input, target = train_provider.next()
        
        # decay learning rate
        if cfg.TRAIN.end_lr == cfg.TRAIN.base_lr:
            current_lr = cfg.TRAIN.base_lr
        else:
            current_lr = calculate_lr(iters)
            for param_group in optimizer.param_groups:
                param_group['lr'] = current_lr
        
        optimizer.zero_grad()
        input = F.pad(input, (PAD, PAD, PAD, PAD))
        model.train()
        pred = model(input)
        # if not cfg.MODEL.loss_func_logits:
        #     for i in range(len(pred)):
        #         pred[i] = torch.sigmoid(pred[i])
        pred_depad = []
        for p in pred:
            pred_depad.append(F.pad(p, (-PAD, -PAD, -PAD, -PAD)))
        del pred
        # target = torch.unsqueeze(target, dim=1)
        # if iters == 1:
        #     writer.add_graph(model, (input,))

        ############################## Compute Loss ########################################
        # if cfg.MODEL.loss_balance_weight:
        #     cur_weight = edge_weight(target)
        #     writer.add_histogram('weight_edge', cur_weight.clone().cpu().data.numpy(), iters)
        # else:
        #     cur_weight = None

        loss1, loss2, loss3 = multiloss(pred_depad, target)
        pred_depad_2 = pred_depad[2][:,1]
        pred_depad_2 = torch.unsqueeze(pred_depad_2, dim=1)
        pred_ = torch.cat([pred_depad_2, pred_depad_2, pred_depad_2], dim=1)
        target_ = torch.cat([target, target, target], dim=1)
        # loss_twonet = 0.25*loss1+loss2+loss3
        loss_twonet = 0.75*(0.5*loss1+loss2)+loss3
        # loss_vgg2 = loss_vgg(pred_, target_)
        out1 = model_vgg(pred_)
        out2 = model_vgg(target_)
        loss_vgg2 = vgg_loss(out1, out2)
        loss = loss_twonet + vgg_weight * loss_vgg2
        
        loss.backward()
        if cfg.TRAIN.weight_decay is not None:
            for group in optimizer.param_groups:
                for param in group['params']:
                    param.data = param.data.add(-cfg.TRAIN.weight_decay * group['lr'], param.data)
        optimizer.step()
        
        sum_loss += loss.item()
        sum_time += time.time() - t1
        
        # log train
        if iters % cfg.TRAIN.display_freq == 0:
            rcd_time.append(sum_time)
            logging.info('step %d, loss = %.4f (wt: *10, lr: %.8f, et: %.2f sec, rd: %.2f min)'
                            % (iters, sum_loss / cfg.TRAIN.display_freq * 10, current_lr, sum_time,
                            (cfg.TRAIN.total_iters - iters) / cfg.TRAIN.display_freq * np.mean(np.asarray(rcd_time)) / 60))
            logging.info('step %d, loss_twonet = %.4f, loss_vgg = %.4f' % (iters, loss_twonet, vgg_weight * loss_vgg2))
            writer.add_scalar('loss/loss1', loss1, iters)
            writer.add_scalar('loss/loss2', loss2, iters)
            writer.add_scalar('loss/loss3', loss3, iters)
            writer.add_scalar('loss/loss_vgg', loss_vgg2, iters)
            writer.add_scalar('final_loss', sum_loss / cfg.TRAIN.display_freq * 10, iters)
            f_loss_txt.write('step = ' + str(iters) + ', loss = ' + str(sum_loss / cfg.TRAIN.display_freq * 10))
            f_loss_txt.write('\n')
            f_loss_txt.flush()
            sys.stdout.flush()
            sum_time = 0
            sum_loss = 0
        
        # valid
        if iters % cfg.TRAIN.valid_freq == 0:
            input = F.pad(input, (-PAD, -PAD, -PAD, -PAD))
            input0 = (np.squeeze(input[0].data.cpu().numpy()) * 255).astype(np.uint8)
            target = (np.squeeze(target[0].data.cpu().numpy()) * 255).astype(np.uint8)
            pred_show = []
            for p in pred_depad:
                temp = np.squeeze(p[0].data.cpu().numpy())
                temp[temp>1] = 1; temp[temp<0] = 0
                temp = (temp * 255).astype(np.uint8)
                pred_show.append(temp)
            # input0 = input0[0]
            white = np.ones_like(pred_show[2][1], dtype=np.uint8)
            im1 = np.concatenate([input0, pred_show[0][1], pred_show[1][1]], axis=1)
            im2 = np.concatenate([target, pred_show[2][1], white], axis=1)
            im_cat = np.concatenate([im1, im2], axis=0)
            Image.fromarray(im_cat).save(os.path.join(cfg.cache_path, '%06d.png' % iters))
        
        # save
        if iters % cfg.TRAIN.save_freq == 0:
            model.eval()
            for k in range(valid_provider.data.num):
                raw, label = valid_provider.data.gen(k)
                crop_img = Crop_image(raw,crop_size=cfg.TRAIN.crop_size,overlap=cfg.TRAIN.overlap)
                for i in range(crop_img.num):
                    for j in range(crop_img.num):
                        raw_crop = crop_img.gen(i, j)
                        #########
                        #inference
                        if crop_img.dim == 3:
                            raw_crop_ = raw_crop[np.newaxis, :, :, :]
                        else:
                            raw_crop_ = raw_crop[np.newaxis, np.newaxis, :, :]
                        inputs = torch.Tensor(raw_crop_).to(device)
                        inputs = F.pad(inputs, (PAD, PAD, PAD, PAD))
                        
                        with torch.no_grad():
                            pred = model(inputs)
                        pred = pred[2]
                        pred = F.pad(pred, (-PAD, -PAD, -PAD, -PAD))
                        pred = F.softmax(pred, dim=1)
                        pred = torch.argmax(pred, dim=1).squeeze(0)
                        pred = pred.data.cpu().numpy()
                        # pred = np.squeeze(pred)
                        # pred = pred[1]
                        #########
                        crop_img.save(i, j, pred)
                results = crop_img.result()
                results[results<=thresd] = 0
                results[results>thresd] = 1
                temp_label = label.flatten()
                temp_result = results.flatten()
                f1_common = f1_score(1 - temp_label, 1 - temp_result)
                f1_soft = 0
                results_img = (results * 255).astype(np.uint8)
                label_img = (label * 255).astype(np.uint8)
                im_cat_valid = np.concatenate([results_img, label_img], axis=1)
                if k == valid_provider.data.num - 1:
                    Image.fromarray(im_cat_valid).save(os.path.join(cfg.cache_path, 'valid_%06d.png' % iters))
                valid_score.append(f1_soft)
                valid_score_tmp.append(f1_common)
            avg_f1 = sum(valid_score) / len(valid_score)
            avg_f1_tmp = sum(valid_score_tmp) / len(valid_score_tmp)
            logging.info('step %d, f1_soft = %.6f' % (iters, avg_f1))
            logging.info('step %d, f1_common = %.6f' % (iters, avg_f1_tmp))
            writer.add_scalar('valid', avg_f1_tmp, iters)
            f_valid_txt.write('step %d, f1 = %.6f' % (iters, avg_f1_tmp))
            f_valid_txt.write('\n')
            f_valid_txt.flush()
            sys.stdout.flush()
            if avg_f1_tmp > most_f1:
                most_f1 = avg_f1_tmp
                most_iters = iters
                states = {'current_iter': iters, 'valid_result': avg_f1_tmp,
                        'model_weights': model.state_dict()}
                torch.save(states, os.path.join(cfg.save_path, 'model.ckpt'))
                print('***************save modol, when f1 = %.6f and iters = %d.***************' % (most_f1, most_iters))