ref_training.py

import numpy as np
import math
import os
import argparse
from tqdm import tqdm
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

class RefinementDataset(Dataset):
    def __init__(self, device, data_path='.', train=False):
        if train:
            input_para = np.load(os.path.join(data_path, 'param_src_train.npy'))
            target_para = np.load(os.path.join(data_path, 'param_tgt_train.npy'))
            noisy_img = np.load(os.path.join(data_path, 'images_noisy_train.npy'))
            gt_img = np.load(os.path.join(data_path, 'images_gt_train.npy'))
            alpha = np.load(os.path.join(data_path, 'alpha_train.npy'))
        else:
            input_para = np.load(os.path.join(data_path, 'param_src_test.npy'))
            target_para = np.load(os.path.join(data_path, 'param_tgt_test.npy'))
            noisy_img = np.load(os.path.join(data_path, 'images_noisy_test.npy'))
            gt_img = np.load(os.path.join(data_path, 'images_gt_test.npy'))
            alpha = np.load(os.path.join(data_path, 'alpha_test.npy'))
        self.input_para = torch.from_numpy(input_para).float().flatten(start_dim=1,end_dim=2).to(device)
        self.target_para = torch.from_numpy(target_para).float().flatten(start_dim=1,end_dim=2).to(device)
        self.noisy_img = torch.from_numpy(noisy_img).float().to(device)
        self.gt_img = torch.from_numpy(gt_img).float().to(device)
        self.alpha = torch.from_numpy(alpha).float().to(device)
    def __len__(self):
        return self.input_para.shape[0]
    def __getitem__(self, idx):
        return self.input_para[idx, :, :], self.target_para[idx, :, :], self.noisy_img[idx, :, :, :], self.gt_img[idx, :, :, :], self.alpha[idx]

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len, stride, device = None):
        super(PositionalEncoding, self).__init__()
        d_model_half = int(d_model / 2)
        position = torch.linspace(0, (max_len-1)*stride, max_len)
        pe = torch.zeros((max_len, max_len, d_model))
        div_term = torch.exp(torch.arange(0, d_model_half, 2) * (-2 * math.log(10000.0) / d_model)).unsqueeze(0).unsqueeze(0)
        pe[:, :, 0:d_model_half:2] = torch.sin(position.unsqueeze(1).unsqueeze(1) * div_term)
        pe[:, :, 1:d_model_half:2] = torch.cos(position.unsqueeze(1).unsqueeze(1) * div_term)
        pe[:, :, d_model_half:d_model:2] = torch.sin(position.unsqueeze(0).unsqueeze(2) * div_term)
        pe[:, :, d_model_half+1:d_model:2] = torch.cos(position.unsqueeze(0).unsqueeze(2) * div_term)
        self.pe = pe.flatten(start_dim=0, end_dim=1).unsqueeze(0).to(device)
    def forward(self, x):
        x += self.pe[:, :x.size(1), :]
        return x

class TransformerRefinement(nn.Module):
    def __init__(self, max_len=64, stride=2, in_parameter_size=14, out_parameter_size=5, d_model=128, nhead=8, num_encoder_layers=8, num_decoder_layers=0, \
                    dim_feedforward=256, batch_first=True, device=None):
        super(TransformerRefinement, self).__init__()
        self.in_src_projection = nn.Linear(in_features=in_parameter_size, out_features=d_model)
        self.positional_encoding = PositionalEncoding(d_model=d_model, max_len=max_len, stride=stride, device=device)
        self.transformer = nn.Transformer(d_model=d_model, nhead=nhead, num_encoder_layers=num_encoder_layers, num_decoder_layers=num_decoder_layers, \
                    dim_feedforward=dim_feedforward, batch_first=batch_first)
        self.generator = nn.Linear(in_features=d_model, out_features=out_parameter_size)
    def forward(self, src: torch.Tensor):
        src_emb = self.positional_encoding(self.in_src_projection(src))
        outs = self.transformer.encoder(src_emb)
        outs = self.generator(outs)
        return outs

def evaluateDataset(args, model, criteria, datasetloader, data_size, init_train=True):
    model.eval()
    with torch.no_grad():
        total_loss = 0
        for src, tgt, ny_img, gt_img, alpha in datasetloader:
            est = model(src)
            if not init_train:
                loss = criteria(est, ny_img, gt_img, alpha)
            else:
                loss = criteria(est, tgt)
            total_loss += loss
        num_batch = data_size // args.batch_size
        avg_total_loss = total_loss / num_batch
        return avg_total_loss

def showCurve(args, points, figname):
    plt.figure(figsize = (8,6))
    plt.xlabel('Epochs')
    plt.ylabel('Average loss')
    epoches_num = np.arange(points.shape[0])
    plt.yscale("log")
    plt.plot(epoches_num, points, linestyle='-', color='b', linewidth=2)
    cf = plt.gcf()
    cf.savefig('%s%s.jpg'%(args.data_path, figname), format='jpg', bbox_inches='tight', dpi=600)

class L2Loss(nn.Module):
    def __init__(self, R, batch_size, stride, eta, delta, lmbda_boundary, lmbda_color, device):
        super(L2Loss, self).__init__()
        y, x = torch.meshgrid([torch.linspace(-1.0, 1.0, R), \
                                torch.linspace(-1.0, 1.0, R)])
        self.x = x.view(1, R, R, 1, 1).to(device)
        self.y = y.view(1, R, R, 1, 1).to(device)
        self.R = R
        self.batch_size = batch_size
        self.eta = eta
        self.delta = delta
        self.stride = stride
        self.lmbda_boundary = lmbda_boundary
        self.lmbda_color = lmbda_color
        self.H = 147
        self.W = 147
        self.H_patches = 64
        self.W_patches = 64
        self.num_patches = torch.nn.Fold(output_size=[self.H, self.W],
                                         kernel_size=R,
                                         stride=stride)(torch.ones(1, R**2,
                                                                    self.H_patches * self.W_patches,
                                                                    device=device)).view(self.H, self.W)
    
    def params2dists(self, params, tau=1e-1):
        x0     = params[:, 3, :, :].unsqueeze(1).unsqueeze(1)
        y0     = params[:, 4, :, :].unsqueeze(1).unsqueeze(1)

        angles = torch.remainder(params[:, :3, :, :], 2 * np.pi)
        angles = torch.sort(angles, dim=1)[0]
        angle1 = angles[:, 0, :, :].unsqueeze(1).unsqueeze(1)
        angle2 = angles[:, 1, :, :].unsqueeze(1).unsqueeze(1)
        angle3 = angles[:, 2, :, :].unsqueeze(1).unsqueeze(1)
        angle4 = 0.5 * (angle1 + angle3) + \
                     torch.where(torch.remainder(0.5 * (angle1 - angle3), 2 * np.pi) >= np.pi,
                                 torch.ones_like(angle1) * np.pi, torch.zeros_like(angle1))
        def g(dtheta):
            return (dtheta / np.pi - 1.0) ** 35
        
        sgn42 = torch.where(torch.remainder(angle2 - angle4, 2 * np.pi) < np.pi,
                            torch.ones_like(angle2), -torch.ones_like(angle2))
        tau42 = g(torch.remainder(angle2 - angle4, 2*np.pi)) * tau
        dist42 = sgn42 * torch.min( sgn42 * (-torch.sin(angle4) * (self.x - x0) + torch.cos(angle4) * (self.y - y0)),
                                   -sgn42 * (-torch.sin(angle2) * (self.x - x0) + torch.cos(angle2) * (self.y - y0))) + tau42
        
        sgn13 = torch.where(torch.remainder(angle3 - angle1, 2 * np.pi) < np.pi,
                            torch.ones_like(angle3), -torch.ones_like(angle3))
        tau13 = g(torch.remainder(angle3 - angle1, 2*np.pi)) * tau
        dist13 = sgn13 * torch.min( sgn13 * (-torch.sin(angle1) * (self.x - x0) + torch.cos(angle1) * (self.y - y0)),
                                   -sgn13 * (-torch.sin(angle3) * (self.x - x0) + torch.cos(angle3) * (self.y - y0))) + tau13
        return torch.stack([dist13, dist42], dim=1)

    def dists2indicators(self, dists):
        hdists = 0.5 * (1.0 + (2.0 / np.pi) * torch.atan(dists / self.eta))
        return torch.stack([1.0 - hdists[:, 0, :, :, :, :],
                                  hdists[:, 0, :, :, :, :] * (1.0 - hdists[:, 1, :, :, :, :]),
                                  hdists[:, 0, :, :, :, :] *        hdists[:, 1, :, :, :, :]], dim=1)
    
    def get_dists_and_patches(self, params):
        dists = self.params2dists(params)
        wedges = self.dists2indicators(dists)
        colors = (self.img_patches.unsqueeze(2) * wedges.unsqueeze(1)).sum(-3).sum(-3) / \
                    (wedges.sum(-3).sum(-3).unsqueeze(1) + 1e-10)
        patches = (wedges.unsqueeze(1) * colors.unsqueeze(-3).unsqueeze(-3)).sum(dim=2)
        return dists, colors, patches

    def local2global(self, patches):
        N = patches.shape[0]
        C = patches.shape[1]
        return torch.nn.Fold(output_size=[self.H, self.W], kernel_size=self.R, stride=self.stride)(
            patches.view(N, C*self.R**2, -1)).view(N, C, self.H, self.W) / \
                self.num_patches.unsqueeze(0).unsqueeze(0)

    def dists2boundaries(self, dists):
        d1 = dists[:, 0:1, :, :, :, :]
        d2 = dists[:, 1:2, :, :, :, :]
        minabsdist = torch.where(d1 < 0.0, -d1, torch.where(d2 < 0.0, torch.min(d1, -d2), torch.min(d1, d2)))
        return 1.0 / (1.0 + (minabsdist / self.delta) ** 2)

    def get_boundary_consistency_term(self, dists):
        curr_global_boundaries_patches = nn.Unfold(self.R, stride=self.stride)(
            self.global_boundaries.detach()).view(self.batch_size, 1, self.R, self.R, self.H_patches, self.W_patches)
        local_boundaries = self.dists2boundaries(dists)
        consistency = ((local_boundaries - curr_global_boundaries_patches) ** 2).mean(-3).mean(-3)
        return consistency
    
    def get_color_consistency_term(self, dists, colors, alpha):
        curr_global_image_patches = nn.Unfold(self.R, stride=self.stride)(
            self.global_image.detach()).view(self.batch_size, 3, self.R, self.R, self.H_patches, self.W_patches)
        wedges = self.dists2indicators(dists)
        consistency = (wedges.unsqueeze(1) * (
            colors.unsqueeze(-3).unsqueeze(-3) - curr_global_image_patches.unsqueeze(2)) ** 2).mean(-3).mean(-3).sum(1).sum(1) / alpha[:, None, None]
        return consistency
    
    def get_loss(self, dists, colors, patches, alpha):
        loss_per_patch = ((self.gt_img_patches - patches) ** 2).mean(-3).mean(-3).sum(1) / alpha[:, None, None] + \
                            self.lmbda_boundary * self.get_boundary_consistency_term(dists) + \
                            self.lmbda_color * self.get_color_consistency_term(dists, colors, alpha)
        return loss_per_patch
    
    def forward(self, ests, noisy_image, gt_image, alpha):
        ests = ests.permute(0,2,1).view(self.batch_size, 5, self.H_patches, self.W_patches)
        angles = (ests[:, :3, :, :] + 1) * np.pi
        x0y0 = ests[:, 3:, :, :] * 3
        para = torch.cat([angles, x0y0], dim=1)
        self.img_patches = nn.Unfold(self.R, stride=self.stride)(noisy_image.permute(0,3,1,2)).view(self.batch_size, 3, self.R, self.R, self.H_patches, self.W_patches)
        self.gt_img_patches = nn.Unfold(self.R, stride=self.stride)(gt_image.permute(0,3,1,2)).view(self.batch_size, 3, self.R, self.R, self.H_patches, self.W_patches) #!
        dists, colors, patches = self.get_dists_and_patches(para)
        self.global_image = self.local2global(patches)
        self.global_boundaries = self.local2global(self.dists2boundaries(dists))
        loss = self.get_loss(dists, colors, patches, alpha).mean()
        return loss

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--cuda', type=str, default='cuda:0', help='Enable cuda')
    parser.add_argument('--epoch_num', type=int, default=1700, help='Number of epoches')
    parser.add_argument('--learning_rate_1', type=float, default=5e-5, help='Learning rate for initial training')
    parser.add_argument('--learning_rate_2', type=float, default=3.5e-4, help='Initial learning rate for late training')
    parser.add_argument('--loss_change', type=int, default=100, help='Epoch to change the loss function')
    parser.add_argument('--batch_size', type=int, default=16, help='Number of batch size')
    parser.add_argument('--R', type=int, default=21, help='Patch size')
    parser.add_argument('--stride', type=int, default=2, help='Stride')
    parser.add_argument('--lambda_boundary', type=float, default=0.5, help='boundary weight in loss function')
    parser.add_argument('--lambda_color', type=float, default=0.1, help='color weight in loss function')
    parser.add_argument('--data_path', type=str, default='./dataset/refinement/', help='Path of dataset')
    parser.add_argument('--input_size', type=int, default=14, help='Input layer size')
    parser.add_argument('--output_size', type=int, default=5, help='Output layer size')
    parser.add_argument('--eta', type=float, default=0.01, help='Width parameter for Heaviside functions')
    parser.add_argument('--delta', type=float, default=0.05, help='delta in loss function')
    args = parser.parse_args()

    np.random.seed(1896)
    torch.manual_seed(1896)
    device = torch.device(args.cuda)

    dataset_train = RefinementDataset(device, data_path=args.data_path, train=True)
    train_loader = DataLoader(dataset_train, batch_size=args.batch_size, shuffle=True, drop_last=True)
    dataset_test = RefinementDataset(device, data_path=args.data_path, train=False)
    test_loader = DataLoader(dataset_test, batch_size=args.batch_size, shuffle=False, drop_last=True)

    refiner = TransformerRefinement(in_parameter_size = args.input_size, out_parameter_size = args.output_size, device=device).to(device)
    for p in refiner.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
    optimizer = torch.optim.Adam(refiner.parameters(), lr=args.learning_rate_1)
    criteria_1 = nn.MSELoss()
    criteria_2 = L2Loss(args.R, args.batch_size, args.stride, args.eta, args.delta, args.lambda_boundary, args.lambda_color, device)

    lr_updated = 0
    best_avg_loss = np.inf
    best_epoch = 0
    avg_total_loss = np.zeros((args.epoch_num,), dtype=float)
    for epoch in tqdm(range(args.epoch_num)):
        refiner.train()
        if epoch == args.loss_change:
            scheduler_2 = torch.optim.lr_scheduler.CyclicLR(optimizer, args.learning_rate_2*0.5, args.learning_rate_2, step_size_up=200, step_size_down=200, mode='triangular2', cycle_momentum=False)
        if epoch > args.loss_change:
            scheduler_2.step()
        for step, (src, tgt, ny_img, gt_img, alpha) in enumerate(train_loader):
            est = refiner(src)
            optimizer.zero_grad()
            if epoch >= args.loss_change:
                loss = criteria_2(est, gt_img, gt_img, alpha)
            else:
                loss = criteria_1(est, tgt)
            loss.backward()
            optimizer.step()
        if epoch >= args.loss_change:
            avg_total_loss[epoch] = evaluateDataset(args, refiner, criteria_2, test_loader, len(dataset_test), init_train=False)
        else:
            avg_total_loss[epoch] = evaluateDataset(args, refiner, criteria_1, test_loader, len(dataset_test), init_train=True)
        if epoch >= args.loss_change:
            if avg_total_loss[epoch] < best_avg_loss:
                best_avg_loss = avg_total_loss[epoch]
                torch.save(refiner.state_dict(), '%sbest_ran.pth'%args.data_path)
                best_epoch = epoch
    np.save('%sloss_curve_total.npy'%args.data_path, avg_total_loss)
    showCurve(args, avg_total_loss[:args.loss_change], 'loss_curve_total_tf_1')
    showCurve(args, avg_total_loss[args.loss_change:], 'loss_curve_total_tf_2')
    print('-- Best epoch is {}, with average loss of {}'.format(best_epoch, best_avg_loss))