Browse Source

scale

master
root 2 months ago
parent
commit
32335107fd
6 changed files with 47 additions and 49 deletions
  1. +2
    -1
      .gitignore
  2. +18
    -9
      eval.py
  3. +17
    -35
      train.py
  4. +4
    -2
      train_ModelArts.py
  5. +3
    -1
      train_psnr.py
  6. +3
    -1
      train_psnr_ModelArts.py

+ 2
- 1
.gitignore View File

@@ -113,4 +113,5 @@ analyze_fail.dat
src/model/vgg19_ImageNet.ckpt
DIV2K.zip
src/model/psnr-1_31523.ckpt
analyze_fail.dat
analyze_fail.dat
images/

+ 18
- 9
eval.py View File

@@ -7,10 +7,8 @@ import glob
import numpy as np
import cv2
import mindspore.nn as nn
from PIL import Image
from mindspore.nn import PSNR,SSIM
from mindspore import Tensor, context
from mindspore.nn.optim.momentum import Momentum
from mindspore.train.model import Model
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore.ops import operations as P
@@ -74,16 +72,16 @@ def test():
nf=config_psnr["G_nf"],
nb=config_psnr["G_nb"],
)
dataset = get_dataset_DIV2K(
dataset,dataset_len = get_dataset_DIV2K(
base_dir="./data",
downsample_factor=config_psnr["down_factor"],
mode="eval",
mode="valid",
aug=False,
repeat=1,
num_readers=4,
shard_id=args.rank,
shard_num=args.group_size,
batch_size=args.batch_size,
batch_size=1,
)

eval_net = BuildEvalNetwork(model_psnr)
@@ -94,8 +92,11 @@ def test():
eval_net.set_train(False)
ssim = nn.SSIM()
psnr = nn.PSNR()
test_data_iter = dataset.create_dict_iter(out_numpy)

test_data_iter = dataset.create_dict_iter(out_numpy=False)
psnr_bic_all = 0.0
psnr_real_all = 0.0
ssim_bic_all = 0.0
ssim_real_all = 0.0
for i, sample in enumerate(test_data):
lr = sample['inputs']
real_hr = sample['target']
@@ -106,9 +107,17 @@ def test():
psnr_real = psnr(gen_hr,real_hr)
ssim_bic = ssim(gen_hr,bic_hr)
ssim_real = ssim(gen_hr,real_hr)
psnr_bic_all += psnr_bic
psnr_real_all = psnr_real
ssim_bic_all = ssim_bic
ssim_real_all = ssim_real
print(psnr_bic,psnr_real,ssim_bic,ssim_real)
result_img_path = os.path.join(args_opt.results_path + "DIV2K", 'Bic_SR_HR_' + str(i))
if i%10 == 0:
results_img = np.concatenate((bic_img.asnumpy(), sr_img.asnumpy(), hr_img.asnumpy()), 1)
results_img = np.concatenate((bic_img[0].asnumpy(), sr_img[0].asnumpy(), hr_img[0].asnumpy()), 1)
cv2.imwrite(result_img_path, results_img)
psnr_bic_all += psnr_bic_all/dataset_len
psnr_real_all = psnr_real_all/dataset_len
ssim_bic_all = ssim_bic_all/dataset_len
ssim_real_all = ssim_real_all/dataset_len
print(psnr_bic_all,psnr_real_all,ssim_bic_all,ssim_real_all)

+ 17
- 35
train.py View File

@@ -3,7 +3,7 @@ import os
import argparse
import ast
import numpy as np
from PIL import Image
import cv2
import mindspore
import mindspore.common.dtype as mstype
from mindspore import nn
@@ -19,6 +19,7 @@ from mindspore.train.callback import (
_InternalCallbackParam,
RunContext,
)
from mindspore.ops import composite as C

from src.model.RRDB_Net import RRDBNet
from src.model.discriminator_net import VGGStyleDiscriminator128
@@ -26,34 +27,6 @@ from src.model.cell import GeneratorLossCell, DiscriminatorLossCell, TrainOneSte
from src.config.config import ESRGAN_config
from src.dataset.dataset_DIV2K import get_dataset_DIV2K

# save image
def save_image(img, img_path):
mul = ops.Mul()
add = ops.Add()
if isinstance(img, Tensor):
img = mul(img, 0.5)
img = add(img, 0.5)
img = img.asnumpy().astype(np.uint8).transpose((0, 2, 3, 1))

elif not isinstance(img, np.ndarray):
raise ValueError("img should be Tensor or numpy array, but get {}".format(type(img)))

IMAGE_SIZE = 64 # Image size
IMAGE_ROW = 8 # Row num
IMAGE_COLUMN = 8 # Column num
PADDING = 2 #Interval of small pictures
to_image = Image.new('RGB', (IMAGE_COLUMN * IMAGE_SIZE + PADDING * (IMAGE_COLUMN + 1),
IMAGE_ROW * IMAGE_SIZE + PADDING * (IMAGE_ROW + 1))) # create a new picture
# cycle
i = 0
for y in range(1, IMAGE_ROW + 1):
for x in range(1, IMAGE_COLUMN + 1):
from_image = Image.fromarray(img[i])
to_image.paste(from_image, ((x - 1) * IMAGE_SIZE + PADDING * x, (y - 1) * IMAGE_SIZE + PADDING * y))
i = i + 1

to_image.save(img_path) #save


def parse_args():
parser = argparse.ArgumentParser("ESRGAN")
@@ -64,6 +37,8 @@ def parse_args():
parser.add_argument(
"--aug", type=bool, default=True, help="Use augement for dataset"
)
parser.add_argument("--loss_scale", type=float,
default=1024.0, help="loss scale")
parser.add_argument('--data_dir', type=str,
default=None, help='Dataset path')
parser.add_argument("--batch_size", type=int, default=16, help="batch_size")
@@ -91,6 +66,13 @@ def parse_args():
return args


# save image
def save_img(img, img_name,save_dir):
save_img = C.clip_by_value(img.squeeze(), 0, 1).asnumpy().transpose(1, 2, 0)
# save img
save_fn = save_dir + '/' + img_name
cv2.imwrite(save_fn, cv2.cvtColor(save_img * 255, cv2.COLOR_BGR2RGB), [cv2.IMWRITE_PNG_COMPRESSION, 0])

def train():
args_opt = parse_args()
config = ESRGAN_config
@@ -135,9 +117,9 @@ def train():
milestone=config["lr_steps"], learning_rates=config["lr_D"]
)
optimizerD = nn.Adam(discriminator.trainable_params(
), learning_rate=lr_D, beta1=0.5, beta2=0.999)
), learning_rate=lr_D, beta1=0.5, beta2=0.999,loss_scale=args_opt.loss_scale)
optimizerG = nn.Adam(generator.trainable_params(
), learning_rate=lr_G, beta1=0.5, beta2=0.999)
), learning_rate=lr_G, beta1=0.5, beta2=0.999,loss_scale=args_opt.loss_scale)

# Define One step train
G_trainOneStep = TrainOneStepCellGen(G_loss_cell, optimizerG)
@@ -221,10 +203,10 @@ def train():
dis_iterations += 1

print('[%d/%d][%d/%d][%d] Loss_D: %10f Loss_G: %10f'
% (epoch, args_opt.epoch_size, i, length, dis_iterations,
% (epoch, args_opt.epoch_size, i, length//args_opt.batch_size, dis_iterations,
np.sum(discriminator_loss.asnumpy()), generator_loss.asnumpy()))
if dis_iterations % 10 == 0:
save_image(target, '{0}/real_samples.png'.format(args_opt.experiment))
save_image(fake_hr, '{0}/fake_samples_{1}.png'.format(args_opt.experiment, dis_iterations))
if dis_iterations % 5 == 0:
save_img(target[0], 'real_samples_{0}.png'.format(dis_iterations),args_opt.experiment)
save_img(fake_hr[0], 'fake_samples_{0}.png'.format(dis_iterations),args_opt.experiment)
if __name__ == "__main__":
train()

+ 4
- 2
train_ModelArts.py View File

@@ -68,6 +68,8 @@ def parse_args():
parser.add_argument(
"--aug", type=bool, default=True, help="Use augement for dataset"
)
parser.add_argument("--loss_scale", type=float,
default=1024.0, help="loss scale")
parser.add_argument('--data_dir', type=str,
default=None, help='Dataset path')
parser.add_argument("--batch_size", type=int, default=16, help="batch_size")
@@ -168,9 +170,9 @@ def train():
milestone=config["lr_steps"], learning_rates=config["lr_D"]
)
optimizerD = nn.Adam(discriminator.trainable_params(
), learning_rate=lr_D, beta1=0.5, beta2=0.999)
), learning_rate=lr_D, beta1=0.5, beta2=0.999,loss_scale=args_opt.loss_scale)
optimizerG = nn.Adam(generator.trainable_params(
), learning_rate=lr_G, beta1=0.5, beta2=0.999)
), learning_rate=lr_G, beta1=0.5, beta2=0.999,loss_scale=args_opt.loss_scale)

# Define One step train
G_trainOneStep = TrainOneStepCellGen(G_loss_cell, optimizerG)


+ 3
- 1
train_psnr.py View File

@@ -31,6 +31,8 @@ def parse_args():
default=0,
help="device id of GPU or Ascend. (Default: None)",
)
parser.add_argument("--loss_scale", type=float,
default=1024.0, help="loss scale")
parser.add_argument(
"--aug", type=bool, default=True, help="Use augement for dataset"
)
@@ -77,7 +79,7 @@ def train(config):
milestone=config_psnr["lr_steps"], learning_rates=config_psnr["lr"]
)
opt = nn.Adam(
params=model_psnr.trainable_params(), learning_rate=lr, beta1=0.9, beta2=0.99
params=model_psnr.trainable_params(), learning_rate=lr, beta1=0.9, beta2=0.99,loss_scale=args_opt.loss_scale
)
loss = nn.L1Loss()
loss.add_flags_recursive(fp32=True)


+ 3
- 1
train_psnr_ModelArts.py View File

@@ -39,6 +39,8 @@ def parse_args():
parser.add_argument(
"--aug", type=bool, default=True, help="Use augement for dataset"
)
parser.add_argument("--loss_scale", type=float,
default=1024.0, help="loss scale")
parser.add_argument("--batch_size", type=int, default=4, help="batch_size")
parser.add_argument("--epoch_size", type=int, default=20, help="epoch_size")
parser.add_argument("--rank", type=int, default=0, help="local rank of distributed")
@@ -130,7 +132,7 @@ def train(config):
milestone=config_psnr["lr_steps"], learning_rates=config_psnr["lr"]
)
opt = nn.Adam(
params=model_psnr.trainable_params(), learning_rate=lr, beta1=0.9, beta2=0.99
params=model_psnr.trainable_params(), learning_rate=lr, beta1=0.9, beta2=0.99,loss_scale=args_opt.loss_scale
)
loss = nn.L1Loss()
loss.add_flags_recursive(fp32=True)


Loading…
Cancel
Save