Spaces:
Running
on
T4
Running
on
T4
File size: 3,720 Bytes
561c629 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
# -*- coding: utf-8 -*-
import sys
import os
import torch
# Import files from the local folder
root_path = os.path.abspath('.')
sys.path.append(root_path)
from opt import opt
from architecture.grl import GRL # This place need to adjust for different models
from train_code.train_master import train_master
# Mixed precision training
scaler = torch.cuda.amp.GradScaler()
class train_grl(train_master):
def __init__(self, options, args) -> None:
super().__init__(options, args, "grl") # Pass a model name unique code
def loss_init(self):
# Prepare pixel loss
self.pixel_loss_load()
def call_model(self):
patch_size = 144
window_size = 8
if opt['model_size'] == "small":
# GRL small model
self.generator = GRL(
upscale = opt['scale'],
img_size = patch_size,
window_size = 8,
depths = [4, 4, 4, 4],
embed_dim = 128,
num_heads_window = [2, 2, 2, 2],
num_heads_stripe = [2, 2, 2, 2],
mlp_ratio = 2,
qkv_proj_type = "linear",
anchor_proj_type = "avgpool",
anchor_window_down_factor = 2,
out_proj_type = "linear",
conv_type = "1conv",
upsampler = "pixelshuffle",
).cuda()
elif opt['model_size'] == "tiny":
# GRL tiny model
self.generator = GRL(
upscale = opt['scale'],
img_size = 64,
window_size = 8,
depths = [4, 4, 4, 4],
embed_dim = 64,
num_heads_window = [2, 2, 2, 2],
num_heads_stripe = [2, 2, 2, 2],
mlp_ratio = 2,
qkv_proj_type = "linear",
anchor_proj_type = "avgpool",
anchor_window_down_factor = 2,
out_proj_type = "linear",
conv_type = "1conv",
upsampler = "pixelshuffledirect",
).cuda()
elif opt['model_size'] == "tiny2":
# GRL tiny model
self.generator = GRL(
upscale = opt['scale'],
img_size = 64,
window_size = 8,
depths = [4, 4, 4, 4],
embed_dim = 64,
num_heads_window = [2, 2, 2, 2],
num_heads_stripe = [2, 2, 2, 2],
mlp_ratio = 2,
qkv_proj_type = "linear",
anchor_proj_type = "avgpool",
anchor_window_down_factor = 2,
out_proj_type = "linear",
conv_type = "1conv",
upsampler = "nearest+conv", # Change
).cuda()
else:
raise NotImplementedError("We don't support such model size in GRL model")
# self.generator = torch.compile(self.generator).cuda() # Don't use this for 3090Ti
self.generator.train()
def run(self):
self.master_run()
def calculate_loss(self, gen_hr, imgs_hr):
# Define the loss function here
# Generator pixel loss (l1 loss): generated vs. GT
l_g_pix = self.cri_pix(gen_hr, imgs_hr, self.batch_idx)
self.weight_store["pixel_loss"] = l_g_pix
self.generator_loss += l_g_pix
def tensorboard_report(self, iteration):
# self.writer.add_scalar('Loss/train-Generator_Loss-Iteration', self.generator_loss, iteration)
self.writer.add_scalar('Loss/train-Pixel_Loss-Iteration', self.weight_store["pixel_loss"], iteration)
|