train.py

import os
import math
import argparse
from datetime import datetime
from utils.ssim_pyrotch import SSIM
from torch.utils.data import DataLoader
from datasets import llf_dataset, llf_dataset_all, llf_dataset_random_all, llf_dataset_gain,llf_dataset_all_N, llf_dataset_random_all_N
from model.models import DynamicLFAmp, TestNet
from model.models_api import *
import utils.testHelper as testHelper
import lpips
import logging
import visdom
#from einops import rearrange
log = logging.getLogger()
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
torch.backends.cudnn.benchmark = True


def lf_sheer(lf, disparity):
    # warp LF with one specific disparity
    [batch, height_view, width_view, height, width, channel] = lf.shape

    lf_t = lf.reshape((-1, height, width, channel)).permute((0, 3, 1, 2))  # batch*u*v,c,h,w

    center_u = height_view // 2
    center_v = width_view // 2
    grid = []
    hh = torch.arange(0, height).view(1, height, 1).expand(batch, height, width) # batch h w
    ww = torch.arange(0, width).view(1, 1, width).expand(batch, height, width)
    for u in range(height_view):
        for v in range(width_view):
            dispmap_u = -disparity * (u - center_u)
            dispmap_v = -disparity * (v - center_v)
            h_range = hh + dispmap_u
            w_range = ww + dispmap_v
            # rerange the value to [-1,1]
            h_range = 2. * h_range / (height - 1) - 1
            w_range = 2. * w_range / (width - 1) - 1
            grid_t = torch.stack((w_range, h_range), dim=3)  # [batch,h,w,2]
            grid.append(grid_t)
    grid = torch.cat(grid, 0)  # [batch*u*v,h,w,2]

    warped_lf = F.grid_sample(lf_t, grid.type_as(lf_t), mode='bicubic', padding_mode="zeros", align_corners=True)
    warped_lf = warped_lf.reshape((batch, height_view, width_view, channel, height, width)).permute(
        (0, 1, 2, 4, 5, 3))

    return warped_lf

def get_sheered_lf_volume(lf, disparity_list):
    # warp LF with one specific disparity
    [batch, height_view, width_view, height, width, channel] = lf.shape

    sheered_LF = torch.zeros((batch, height_view, width_view, height, width, channel, len(disparity_list))).type_as(
        lf)

    for i, disparity in enumerate(disparity_list):
        sheered_LF[:, :, :, :, :, :, i] = lf_sheer(lf, disparity)
    # b u v h w c s
    return sheered_LF
class BuildCost(nn.Module):
    def __init__(self, angRes, mindisp, maxdisp):
        super(BuildCost, self).__init__()
        self.angRes = angRes
        self.mindisp = mindisp
        self.maxdisp = maxdisp
        self.oacc = ModulateConv2d(kernel_size=angRes, stride=1)

    def forward(self, x):
        b, c, n, h, w = x.shape
        x = x.permute(0,2,1,3,4).reshape(-1,c,h,w)
        #x = rearrange(x, 'b c (a1 a2) h w -> (b a1 a2) c h w', a1=self.angRes, a2=self.angRes)
        bdr = (self.angRes // 2) * self.maxdisp
        pad = nn.ZeroPad2d((bdr, bdr, bdr, bdr))
        x_pad = pad(x)
        b_p,c_p,h_p,w_p = x_pad.shape
        x_pad = x_pad.reshape(-1,self.angRes,self.angRes,c_p,h_p,w_p).permute(0,3,1,4,2,5).reshape(-1,c_p,self.angRes*h_p,self.angRes*w_p)
        #x_pad = rearrange(x_pad, '(b a1 a2) c h w -> b c (a1 h) (a2 w)', a1=self.angRes, a2=self.angRes)
        h_pad, w_pad = h + 2 * bdr, w + 2 * bdr
        cost = []
        for d in range(self.mindisp, self.maxdisp + 1):
            dila = [h_pad - d, w_pad - d]
            self.oacc.dilation = dila
            crop = (self.angRes // 2) * (d - self.mindisp)
            if d == self.mindisp:
                feat = x_pad
            else:
                feat = x_pad[:, :, crop: -crop, crop: -crop]
            current_cost = self.oacc(feat, h, w)  # b c*n h w
            cost.append(current_cost)
        cost = torch.stack(cost, dim=3)
        return cost

class ModulateConv2d(nn.Module):
    def __init__(self, kernel_size, stride=1, dilation=1):
        super(ModulateConv2d, self).__init__()
        self.kernel_size = kernel_size
        self.stride = stride
        self.dilation = dilation

    def forward(self, x, h, w):
        Unfold = nn.Unfold(kernel_size=self.kernel_size, stride=self.stride, dilation=self.dilation)
        x_unfold = Unfold(x)
        x_unfold_modulated = x_unfold
        Fold = nn.Fold(output_size=(h, w), kernel_size=1, stride=1)
        out = Fold(x_unfold_modulated)
        b_o,_,h_o,w_o = out.shape
        out = out.reshape(b_o,-1,self.kernel_size**2,h_o,w_o)
        #out = rearrange(out, 'b (c n) h w -> b c n h w', n=self.kernel_size**2)
        return out

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True


def parse_args():
    parser = argparse.ArgumentParser()

    parser.add_argument("--data_dir", type=str, default="H:/GcC/lowlightdata")
    # parser.add_argument('-data', "--data_dir", type=str, default="../data")
    parser.add_argument("--output_dir", type=str, default="H:/GcC/outs")
    parser.add_argument("--exp_desp", type=str, default="DYM-sas-var-5-1") # base #shear #Constraint
    # parser.add_argument("--exp_desp", type=str, default="deformX3-epi-mixdata")
    parser.add_argument("--model_type", type=str, default="sas")
    parser.add_argument("--loss_type", type=str, default="var") # base # self+ # vgg+
    parser.add_argument("--dataset", type=str, default="all")
    parser.add_argument("--epochs", type=int, default=10000)
    parser.add_argument("--learning_rate", type=float, default=1e-4)
    parser.add_argument("--gpu_no", type=str, default="0")

    parser.add_argument("--ckpt", type=str, default="")
    parser.add_argument("--patch", type=int, default=128)

    parser.add_argument("--n_view", type=int, default=5)
    parser.add_argument("--is_pai", type=bool, default=False)
    return parser.parse_args()

def lf2epi(lf):
    N, an2, c, h, w = lf.shape
    an = int(math.sqrt(an2))
    # [N,an2,h,w] -> [N*ah*h,aw,w]  &  [N*aw*w,ah,h]
    # print(an)
    # print(lf.view(N,an,an,h,w).permute(0,1,3,2,4).view(-1,an,w).shape)
    #               0  1   2   3  4  5
    epi_h = lf.view(N, an, an, c, h, w).permute(0, 1, 4, 3, 2, 5).contiguous().view(-1, c, an, w)
    epi_v = lf.view(N, an, an, c, h, w).permute(0, 2, 5, 3, 1, 4).contiguous().view(-1, c, an, h)
    return epi_h, epi_v


class GaussianBlurConv(nn.Module):
    def __init__(self, channels=3):
        super(GaussianBlurConv, self).__init__()
        self.channels = channels
        kernel = [[0.00078633, 0.00655965, 0.01330373, 0.00655965, 0.00078633],
                  [0.00655965, 0.05472157, 0.11098164, 0.05472157, 0.00655965],
                  [0.01330373, 0.11098164, 0.22508352, 0.11098164, 0.01330373],
                  [0.00655965, 0.05472157, 0.11098164, 0.05472157, 0.00655965],
                  [0.00078633, 0.00655965, 0.01330373, 0.00655965, 0.00078633]]
        kernel = torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0) # 1， 1 ， 5， 5
        kernel = np.repeat(kernel, self.channels, axis=0)
        self.weight = nn.Parameter(data=kernel, requires_grad=False)
 
    def __call__(self, x):
        # x : [B, c, h, w]
        x = F.conv2d(x, self.weight, padding=2, groups=self.channels)
        return x.permute(0, 2, 3, 1)

class Loss_Fn(nn.Module):
    def __init__(self, an=8,loss_type = None):
        super(Loss_Fn, self).__init__()
        self.an = an
        self.l2 = nn.MSELoss()
        self.l1 = nn.L1Loss()
        self.ssim = SSIM(val_range=1)
        self.epi_ssim = SSIM(window_size=5, val_range=1)
        # self.blur_conv = GaussianBlurConv().cuda()
        # self.lpips_func = lpips.LPIPS(net='vgg').cuda()
        self.blur_conv = GaussianBlurConv()
        self.lpips_func = lpips.LPIPS(net='vgg')
        self.loss_type = loss_type
        self.cost = BuildCost(self.an, -1, 1)


    def uv_varloss_shear(self,r):
        N, an2, c, h, w = r.shape
        # print(r.shape)
        an = int(math.sqrt(an2))
        r = r.permute(0,1,3,4,2).reshape(N,an,an,h,w,c)
        cost = get_sheered_lf_volume(r,[-1,0,1]) #1,u,v,h,w,1,S
        # cost = self.cost(r) #b c n d h w'''
        cost = cost[:,:,:,an:-an,an:-an,:,:].reshape(N,an2,(h-2*an)*(h-2*an),c,3)
        # r = r.permute(0, 1, 3, 4,2).reshape()
        # cost = self.cost(r)  # b c n d h w''
        # cost = cost[:, :, :, :, 2:-2, 2:-2]
        # cost = cost.reshape(N, c, an2, 3, (h - 4) * (w - 4)).permute(0, 4, 3, 1, 2)
        # '
        #cost = cost.reshape(N, c, an2, 3, h * w).permute(0, 4, 3, 1, 2)
        # cost = cost.reshape(-1, c, an2)
        grad = torch.var(cost,dim = 1)
        grad = grad.reshape(N * (h-2*an) * (w-2*an),c, 3)
        grad = torch.min(grad, dim=2)[0]
        grad = torch.sum(grad, dim=1)
        return torch.mean(grad)
    
    def lf2epi(self, lf):
        N, _, c, h, w = lf.shape
        an = self.an
        #               0  1   2    3  4  5
        epi_h = lf.view(N, an, an, c, h, w).permute(0, 1, 4, 3, 2, 5).contiguous().view(-1, c, an, w)
        epi_v = lf.view(N, an, an, c, h, w).permute(0, 2, 5, 3, 1, 4).contiguous().view(-1, c, an, h)

        # epi_h = lf.view(N, an, an, c, h, w).permute(0, 4, 2, 3, 1, 5).contiguous().view(-1, c, an, w)
        # epi_v = lf.view(N, an, an, c, h, w).permute(0, 1, 5, 3, 4, 2).contiguous().view(-1, c, h, an)
        return epi_h, epi_v
    
    def cal_epi_loss(self, x, gt):
        epi_h, epi_v = self.lf2epi(x)
        epi_h_, epi_v_ = self.lf2epi(gt)
        # epi_ssim = self.epi_ssim(epi_h_dx, epi_h_dx_) + self.epi_ssim(epi_h_dy, epi_h_dy_) + \
        #     self.epi_ssim(epi_v_dx, epi_v_dx_) + self.epi_ssim(epi_v_dy, epi_v_dy_)
        epi_ssim = self.epi_ssim(epi_h, epi_h_) + self.epi_ssim(epi_v, epi_v_)
        # epi_ssim = self.lpips_func(epi_h, epi_h_) + self.lpips_func(epi_v, epi_v_)
        
        return (2 - epi_ssim) / 2

    def forward(self, imp,x, gt, r = None):
        # N an2 c h w
        N, an2, c, h, w = x.shape
        pixel_loss = self.l1(x.view(N*an2, c, h, w), gt.view(N*an2, c, h, w))

        ssim_loss = 1 - self.ssim(x.view(N*an2, c, h, w), gt.view(N*an2, c, h, w))
        per_loss = self.lpips_func.forward((r*imp).view(N*an2, c, h, w), gt.view(N*an2, c, h, w))
        tv_loss = 0.0
        tv_loss1 = 0.0
        tv_loss2 = 0.0
        tv_uv_loss = 0.0
        var_loss = 0.0

        if r is not None:
            var_loss = self.uv_varloss_shear(r)


        loss = pixel_loss + 0.5 * ssim_loss + 1 * var_loss

        return loss, pixel_loss,ssim_loss,tv_loss,tv_uv_loss,tv_loss1,tv_loss2, var_loss,per_loss.detach().reshape(N, an2).mean(dim=1)

def overlap_crop_forward(x, model, device, overlap=20, patch_size=128):
    N, an2, c, h, w = x.shape

    output = torch.zeros(N, an2, c, h, w).to(device)

    for i in range(0, h, patch_size):
        for j in range(0, w, patch_size):
            s1, e1 = i, (i + patch_size) if (i + patch_size) < h else h
            s2, e2 = j, (j + patch_size) if (j + patch_size) < w else w
            #
            s11, e11 = s1 - overlap if s1 != 0 else 0, e1 + overlap
            s22, e22 = s2 - overlap if s2 != 0 else 0, e2 + overlap

            patch_input = x[:, :, :, s11:e11, s22:e22]
            patch_input = patch_input.to(device)

            # outs, _ = model(patch_input, hist_data)
            _, patch_out = model(patch_input)
            # with torch.no_grad():
            #     patch_out = model(patch_input)
            # patch_out = outs[-1]

            s1_out = overlap if s1 != 0 else 0
            s2_out = overlap if s2 != 0 else 0

            output[:, :, :, s1:e1, s2:e2] = patch_out[:, :, :, s1_out:s1_out + patch_size,
                                            s2_out:s2_out + patch_size]

    return output


def forward_model(x,r, model):
    res = model(x,r)
    if isinstance(res, tuple):
        return res[-1]
    return res

def forward_model1(x, model, x1):
    res = model(x, x1)
    if isinstance(res, tuple):
        return res[-1]
    return res

def overlap_crop_forward_new(x,r, scale, model, max_length, shave):
    """
    chop for less memory consumption during test
    """
    n_GPUs = 1
    b, u, v, c, h, w = x.size()
    h_half, w_half = h // 2, w // 2
    h_size, w_size = h_half + shave, w_half + shave
    lr_list = [
        x[:, :, :, :, 0:h_size, 0:w_size],
        x[:, :, :, :, 0:h_size, (w - w_size):w],
        x[:, :, :, :, (h - h_size):h, 0:w_size],
        x[:, :, :, :, (h - h_size):h, (w - w_size):w]]
    lr_r_list = [
        r[:, :, :, 0:h_size, 0:w_size],
        r[:, :, :, 0:h_size, (w - w_size):w],
        r[:, :, :, (h - h_size):h, 0:w_size],
        r[:, :, :, (h - h_size):h, (w - w_size):w]]

    sr_list = []
    for i in range(0, 4, n_GPUs):
        lr_batch = torch.cat(lr_list[i:(i + n_GPUs)], dim=0)
        r_batch = torch.cat(lr_r_list[i:(i + n_GPUs)], dim=0)
        if lr_batch.shape[5] > max_length or lr_batch.shape[4] > max_length:
            sr_batch_temp = overlap_crop_forward_new(lr_batch,r_batch, scale, model, max_length, shave)
        else:
            # sr_batch_temp = model(lr_batch.cuda())
            #print(lr_batch.shape)
            #print(r_batch.shape)
            sr_batch_temp = forward_model(lr_batch.cuda(),r_batch.cuda(), model)
            # sr_batch_temp = sr_batch_temp.view(1, u, v, h_size * scale, w_size * scale)

        # sr_batch_temp = model(lr_batch)

        if isinstance(sr_batch_temp, list):
            sr_batch = sr_batch_temp[-1]
        else:
            sr_batch = sr_batch_temp

        sr_list.extend(sr_batch.chunk(n_GPUs, dim=0))

    h, w = scale * h, scale * w
    h_half, w_half = scale * h_half, scale * w_half
    h_size, w_size = scale * h_size, scale * w_size
    shave *= scale

    output = x.new(b, u, v, c, h, w)
    output[:, :, :, :, 0:h_half, 0:w_half] \
        = sr_list[0][:, :, :, :, 0:h_half, 0:w_half]
    output[:, :, :, :, 0:h_half, w_half:w] \
        = sr_list[1][:, :, :, :, 0:h_half, (w_size - w + w_half):w_size]
    output[:, :, :, :, h_half:h, 0:w_half] \
        = sr_list[2][:, :, :, :, (h_size - h + h_half):h_size, 0:w_half]
    output[:, :, :, :, h_half:h, w_half:w] \
        = sr_list[3][:, :, :, :, (h_size - h + h_half):h_size, (w_size - w + w_half):w_size]

    return output



def test_with_same_patch(x, model, bs = 2):
    B, u, v, c, h, w = x.shape
    an2 = u*v
    x = x.reshape(B, an2, c, h, w)
    patches, s1, s2 = testHelper.crop_patch_with_overlap(x)
    n1, n2 = patches.shape[:2]
    res = torch.zeros(patches.shape).cuda()
    patches = patches.view(-1, an2, c, 128, 128)
    res = res.view(-1, an2, c, 128, 128)
    for i in range(n1*n2*B):
        res[i] = model(patches[i].view(1, u, v, c, 128, 128))[-1].view(1, an2, c, 128, 128)
    res = res.view(n1, n2, B, an2, c, 128, 128)
    return testHelper.stitch_patches(x, res, s1, s2)

def test(epoch, device, dataloader_test, model):
    psnr_list, ssim_list, mse_list = [], [], []
    psnr_list1, ssim_list1 = [], []
    ssim = SSIM(val_range=1).cuda()

    l2 = nn.MSELoss()
    t1 = time.time()

    with torch.no_grad():
        for i, (ipt, gt, r) in enumerate(dataloader_test):
            ipt, gt, r = ipt.to(device), gt.to(device), r.to(device)
            N, an2, c, h, w = gt.shape
            an = int(np.sqrt(an2))

            ipt = ipt.reshape(1, an, an, c, h, w)

            out = overlap_crop_forward_new(ipt,r, 1, model, 256, 16)
            # out = test_with_same_patch(ipt, model)
            # out = overlap_crop_forward(input, model, device, 8, 32)
            out = torch.clamp(out, 0, 1)
            out = out.reshape(N, an2, c, h, w)

            for i in range(N):
                view_psnr = []
                view_ssim = []
                for ii in range(an2):
                    mse = l2(gt[i, ii].unsqueeze(0), out[i, ii].unsqueeze(0))
                    psnr = 10 * np.log10(1 / mse.item())
                    view_psnr.append(psnr)
                    ssim_ = ssim(out[i, ii].unsqueeze(0), gt[i, ii].unsqueeze(0)).item()
                    view_ssim.append(ssim_)
                psnr_list.append(np.mean(view_psnr))
                ssim_list.append(np.mean(view_ssim))

            N, an2, c, h, w = gt.shape
            gt = gt.view(N*an2, c, h, w)
            out = out.view(N*an2, c, h, w)

            ssim_ = ssim(out, gt).item()
            mse = l2(gt, out)
            psnr = 10 * np.log10(1 / mse.item())

            psnr_list1.append(psnr)
            ssim_list1.append(ssim_)
            # torch.cuda.empty_cache()
            # print('==>IMG_{} PSNR = {:.4f}, SSIM = {:.4f}'.format(i + 1, psnr, ssim_))

    t2 = time.time()
    avg_psnr, avg_ssim = np.mean(psnr_list), np.mean(ssim_list)
    log.info('===>Test Average PSNR = {:.4f}, SSIM: {:.4f} MSE: {:.4f} Time: {:.2f}'.format(avg_psnr, avg_ssim, 0, t2 - t1))
    log.info('===>Test Average PSNR = {:.4f}, SSIM: {:.4f} MSE: {:.4f} Time: {:.2f}'.format(np.mean(psnr_list1), np.mean(ssim_list1), 0, t2 - t1))
    return avg_psnr, avg_ssim

def train(epoch, device, dataloader, model, optimizer, scheduler, loss_fn=Loss_Fn(an=5)):
    total_loss, count, total_align_loss = 0, 0, 0
    t1 = time.time()
    pixel_loss1, ssim_loss1, tv_loss1, tv_uv_loss1,tvloss2,tvloss3,per_loss1 = 0,0,0,0,0,0,0
    var = 0
    # ssim = SSIM(val_range=1).cuda()
    # loss_fn = Loss_Fn(an=5)

    for i, (ipt, gt, r) in enumerate(dataloader):
        ipt, gt, r = ipt.to(device), gt.to(device), r.to(device)
        #print(r.shape)
        N, an2, c, h, w = gt.shape
        an = int(np.sqrt(an2))
        ipt = ipt.reshape(1, an, an, c, h, w)

        r_lr, out = model(ipt,r) #N, an2, c, h, w
        # out = model(input)
        out = out.reshape(N, an2, c, h, w)
        ipt = ipt.reshape(N, an2, c, h, w)
        loss,pixel_loss,ssim_loss,tv_loss,tv_uv_loss,tvloss22,tvloss32,var_loss,per_loss = loss_fn(ipt,out, gt, r=r_lr)

        total_loss += loss.item()
        pixel_loss1 += pixel_loss.item()
        ssim_loss1 += ssim_loss.item()
        tv_loss1 += tv_loss
        tv_uv_loss1 += tv_uv_loss
        tvloss2 += tvloss22
        tvloss3 += tvloss32
        per_loss1 += per_loss.item()
        var += var_loss
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        count += 1

    if scheduler is not None:
        scheduler.step()

    t2 = time.time()
    log.info('Epoch {:0>4d}: Average loss = {:.3f} pixel_loss1 = {:.3f} ssim_loss1 = {:.3f} tv_loss1= {:.3f}  tv_loss2= {:.3f} tv_loss3= {:.3f} tv_uv_loss1= {:.3f} var_loss1= {:.3f} per_loss1 = {:.3f} per loss = {:.5f} Time: {:.2f}'.format(epoch + 1,
                                                                                         total_loss / count, pixel_loss1/ count,ssim_loss1/ count,tv_loss1/ count, tvloss2/ count,tvloss3/ count,tv_uv_loss1/ count,var/count,per_loss1/ count,
                                                                                         total_align_loss / count, t2 - t1))
    return total_loss / count


def main(args):
    if args.ckpt != "":
        saved_file = os.path.join(args.output_dir, args.exp_desp, args.dataset, args.ckpt)
    else:
        saved_file = os.path.join(args.output_dir, args.exp_desp, args.dataset, datetime.now().strftime("%Y%m%d%H%M"))
        if not os.path.exists(saved_file):
            log.info("created dir:" + saved_file)
            os.makedirs(saved_file)

    print('begin')
    task_name = args.exp_desp + datetime.now().strftime("%Y%m%d%H%M")
    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    
    fh = logging.FileHandler(os.path.join(saved_file, 'Training_%s.log' % datetime.now().strftime("%Y%m%d%H%M")))
    log.addHandler(fh)
    log.info(args)

    if not os.path.exists(saved_file):
        log.info("save dir does not exists!!")
        return 

    device = torch.device("cuda:0")
    model = DeformX3epiWithALTPAdj_S(args.n_view,type_model=args.model_type)


    start_epoch = 0
    best_psnr, best_ssim = 0, 0

    if args.ckpt is not "":
        weights = [i for i in os.listdir(saved_file) if not i.endswith(".log")]
        best_weights = sorted(weights)[-1]
        log.info("selected weight name: {}".format(best_weights))
        model.load_state_dict(torch.load(os.path.join(saved_file, best_weights)))
        tmp = best_weights.split(".p")[0].split("-")
        start_epoch = (int)(tmp[3])
        best_psnr = float(tmp[1])
        best_ssim = float(tmp[2])
    
    model = model.to(device)
    print('model')
    total_num = sum(p.numel() for p in model.parameters())
    log.info("total params = {}".format(total_num))


    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)

    data_train, data_test = llf_dataset_all(args), llf_dataset_random_all(args, is_train=False)
    log.info(len(data_train))
    
    dataloader_train = DataLoader(data_train, batch_size=1, shuffle=True, num_workers=0, pin_memory=False)
    # 测试batch只能为1
    dataloader_test = DataLoader(data_test, batch_size=1, shuffle=False, num_workers=0, pin_memory=False)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2500, gamma=0.5)
    loss_fn = Loss_Fn(an=args.n_view,loss_type=args.loss_type).cuda()
    test_psnr_list = []
    train_loss_list = []
    vis = visdom.Visdom(env=task_name, server='http://172.31.41.137', port=2025)
    for epoch in range(start_epoch, args.epochs):
        train_loss = train(epoch, device, dataloader_train, model, optimizer, scheduler=scheduler, loss_fn=loss_fn)
        train_loss_list.append(train_loss)
        if (epoch) % 10 == 0:
            p, s = test(epoch, device, dataloader_test, model)
            if p > best_psnr:
                best_psnr = p
                best_ssim = s
                torch.save(model.state_dict(), os.path.join(saved_file, 'weights-{:.3f}-{:.3f}-{}.pkl'.format(
                                                                p, s, epoch)))
            log.info('===>Best PSNR: {:.4f} Best SSIM: {:.4f}'.format(best_psnr, best_ssim))
        plot_loss(vis, epoch, train_loss, p)



def plot_loss(vis, epoch, train_loss, test_psnr):
    vis.line(X=np.array([epoch]), Y=np.array([train_loss]), win='loss',
             update='append' if epoch > 0 else None, name='LOSS', opts=dict(title='LOSS'))

    if epoch % 5 == 0:
        vis.line(X=np.array([epoch]), Y=np.array([test_psnr]), win='psnr',
                 update='append' if epoch > 0 else None, name='PSNR', opts=dict(title='PSNR'))

if __name__ == "__main__":
    args = parse_args()
    setup_seed(2021)

    
    main(args)