budui/raycv
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

update

Browse Source
This commit is contained in:
budui 2020-09-05 22:00:17 +08:00
parent 39c754374c
commit e3c760d0c5
12 changed files with 122 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.idea/deployment.xml
View File

@ -1,6 +1,6 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="PublishConfigData" autoUpload="Always" serverName="21d" remoteFilesAllowedToDisappearOnAutoupload="false">
<component name="PublishConfigData" autoUpload="Always" serverName="15d" remoteFilesAllowedToDisappearOnAutoupload="false">
<serverData>
<paths name="14d">
<serverdata>

8
.idea/sshConfigs.xml Normal file
View File

@ -0,0 +1,8 @@
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
<component name="SshConfigs">
<configs>
<sshConfig host="222.20.77.126" id="38d32db7-46b2-4b95-a40c-d17e8eeca6c1" keyPath="C:\Users\wr\.ssh\sg_id_rsa" port="50001" nameFormat="DESCRIPTIVE" username="dancer" />
</configs>
</component>
</project>

10
configs/synthesizers/TAFG.yml
View File

@ -28,7 +28,7 @@ model:
_type: MultiScaleDiscriminator
num_scale: 2
discriminator_cfg:
_type: pix2pixHD
_type: pix2pixHD-PatchDiscriminator
in_channels: 3
base_channels: 64
use_spectral: True
@ -47,10 +47,10 @@ loss:
"11": 0.125
"20": 0.25
"29": 1
criterion: 'NL1'
criterion: 'L2'
style_loss: False
perceptual_loss: True
weight: 5
weight: 0.5
style:
layer_weights:
"1": 0.03125
@ -64,13 +64,13 @@ loss:
weight: 0
fm:
level: 1
weight: 1
weight: 10
recon:
level: 1
weight: 10
style_recon:
level: 1
weight: 10
weight: 0
optimizers:
generator:

34
configs/synthesizers/UGATIT-VoxCeleb2Anime.yml
View File

@ -1,24 +1,20 @@
name: VoxCeleb2Anime
engine: UGATIT
name: selfie2anime-vox2
engine: U-GAT-IT
result_dir: ./result
max_pairs: 1000000
distributed:
model:
# broadcast_buffers: False
misc:
random_seed: 324
handler:
clear_cuda_cache: True
set_epoch_for_dist_sampler: True
checkpoint:
epoch_interval: 1 # one checkpoint every 1 epoch
epoch_interval: 1 # checkpoint once per `epoch_interval` epoch
n_saved: 2
interval:
print_per_iteration: 10 # print once per 10 iteration
tensorboard:
scalar: 10
image: 500
scalar: 100 # log scalar `scalar` times per epoch
image: 2 # log image `image` times per epoch
model:
generator:
@ -27,7 +23,7 @@ model:
out_channels: 3
base_channels: 64
num_blocks: 4
img_size: 128
img_size: 256
light: True
local_discriminator:
_type: UGATIT-Discriminator
@ -74,7 +70,7 @@ data:
target_lr: 0
buffer_size: 50
dataloader:
batch_size: 20
batch_size: 6
shuffle: True
num_workers: 2
pin_memory: True
@ -87,9 +83,9 @@ data:
pipeline:
- Load
- Resize:
size: [ 135, 135 ]
size: [ 286, 286 ]
- RandomCrop:
size: [ 128, 128 ]
size: [ 256, 256 ]
- RandomHorizontalFlip
- ToTensor
- Normalize:
@ -97,7 +93,7 @@ data:
std: [ 0.5, 0.5, 0.5 ]
test:
dataloader:
batch_size: 8
batch_size: 4
shuffle: False
num_workers: 1
pin_memory: False
@ -110,7 +106,7 @@ data:
pipeline:
- Load
- Resize:
size: [ 128, 128 ]
size: [ 256, 256 ]
- ToTensor
- Normalize:
mean: [ 0.5, 0.5, 0.5 ]
@ -122,7 +118,7 @@ data:
pipeline:
- Load
- Resize:
size: [ 128, 128 ]
size: [ 256, 256 ]
- ToTensor
- Normalize:
mean: [ 0.5, 0.5, 0.5 ]

3
configs/synthesizers/UGATIT.yml
View File

@ -120,3 +120,6 @@ data:
- Resize:
size: [ 256, 256 ]
- ToTensor
- Normalize:
mean: [ 0.5, 0.5, 0.5 ]
std: [ 0.5, 0.5, 0.5 ]

8
engine/TAFG.py
View File

@ -90,10 +90,10 @@ class TAFGEngineKernel(EngineKernel):
loss_fm += self.fm_loss(pred_fake[i][j], pred_real[i][j].detach()) / num_scale_discriminator
loss[f"fm_{phase}"] = self.config.loss.fm.weight * loss_fm
loss["recon"] = self.recon_loss(generated["a"], batch["a"]) * self.config.loss.recon.weight
loss["style_recon"] = self.config.loss.style_recon.weight * self.style_recon_loss(
self.generators["main"].module.style_encoders["b"](batch["b"]),
self.generators["main"].module.style_encoders["b"](generated["b"])
)
# loss["style_recon"] = self.config.loss.style_recon.weight * self.style_recon_loss(
# self.generators["main"].module.style_encoders["b"](batch["b"]),
# self.generators["main"].module.style_encoders["b"](generated["b"])
# )
return loss
def criterion_discriminators(self, batch, generated) -> dict:

5
engine/U-GAT-IT.py
View File

@ -148,6 +148,9 @@ class UGATITTestEngineKernel(TestEngineKernel):
def run(task, config, _):
if task == "train":
kernel = UGATITEngineKernel(config)
if task == "test":
run_kernel(task, config, kernel)
elif task == "test":
kernel = UGATITTestEngineKernel(config)
run_kernel(task, config, kernel)
else:
raise NotImplemented

8
engine/base/i2i.py
View File

@ -160,7 +160,7 @@ def get_trainer(config, kernel: EngineKernel):
for k in output["img"]:
image_list = output["img"][k]
tensorboard_handler.writer.add_image(f"train/{k}", make_2d_grid(image_list),
tensorboard_handler.writer.add_image(f"train/{k}", make_2d_grid(image_list, range=(-1, 1)),
engine.state.iteration * pairs_per_iteration)
test_images[k] = []
for i in range(len(image_list)):
@ -183,7 +183,7 @@ def get_trainer(config, kernel: EngineKernel):
for k in test_images:
tensorboard_handler.writer.add_image(
f"test/{k}",
make_2d_grid([torch.cat(ti) for ti in test_images[k]]),
make_2d_grid([torch.cat(ti) for ti in test_images[k]], range=(-1, 1)),
engine.state.iteration * pairs_per_iteration
)
return trainer
@ -218,14 +218,12 @@ def get_tester(config, kernel: TestEngineKernel):
for i in range(batch_size):
image_name = Path(paths[i]).name
torchvision.utils.save_image([img[i] for img in img_tensors], engine.state.img_output_dir / image_name,
nrow=len(img_tensors))
nrow=len(img_tensors), padding=0, normalize=True, range=(-1, 1))
return tester
def run_kernel(task, config, kernel):
assert torch.backends.cudnn.enabled
torch.backends.cudnn.benchmark = True
logger = logging.getLogger(config.name)
with read_write(config):
real_batch_size = config.data.train.dataloader.batch_size * idist.get_world_size()

8
loss/I2I/perceptual_loss.py
View File

@ -98,13 +98,13 @@ class PerceptualLoss(nn.Module):
self.vgg = PerceptualVGG(layer_name_list=list(layer_weights.keys()), vgg_type=vgg_type,
use_input_norm=use_input_norm)
self.criterion = self.set_criterion(criterion)
self.percep_criterion, self.style_criterion = self.set_criterion(criterion)
def set_criterion(self, criterion: str):
assert criterion in ["NL1", "NL2", "L1", "L2"]
norm = F.instance_norm if criterion.startswith("N") else lambda x: x
fn = F.l1_loss if criterion.endswith("L1") else F.mse_loss
return lambda x, t: fn(norm(x), norm(t))
return lambda x, t: fn(norm(x), norm(t)), lambda x, t: fn(x, t)
def forward(self, x, gt):
"""Forward function.
@ -126,7 +126,7 @@ class PerceptualLoss(nn.Module):
if self.perceptual_loss:
percep_loss = 0
for k in x_features.keys():
percep_loss += self.criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
percep_loss += self.percep_criterion(x_features[k], gt_features[k]) * self.layer_weights[k]
else:
percep_loss = None
@ -134,7 +134,7 @@ class PerceptualLoss(nn.Module):
if self.style_loss:
style_loss = 0
for k in x_features.keys():
style_loss += self.criterion(self._gram_mat(x_features[k]), self._gram_mat(gt_features[k])) * \
style_loss += self.style_criterion(self._gram_mat(x_features[k]), self._gram_mat(gt_features[k])) * \
self.layer_weights[k]
else:
style_loss = None

2
main.py
View File

@ -65,6 +65,8 @@ def run(task, config: str, *omega_options, **kwargs):
backup_config = kwargs.get("backup_config", False)
setup_output_dir = kwargs.get("setup_output_dir", False)
setup_random_seed = kwargs.get("setup_random_seed", False)
assert torch.backends.cudnn.enabled
torch.backends.cudnn.benchmark = True
with idist.Parallel(backend=backend) as parallel:
parallel.run(running, conf, task, backup_config=backup_config, setup_output_dir=setup_output_dir,
setup_random_seed=setup_random_seed)

2
model/GAN/UGATIT.py
View File

@ -1,6 +1,6 @@
import torch
import torch.nn as nn
from .residual_generator import ResidualBlock
from .base import ResidualBlock
from model.registry import MODEL

69
tool/verify_loss.py Normal file
View File

@ -0,0 +1,69 @@
import torch
from torch.utils.data import DataLoader
from ignite.utils import convert_tensor
from omegaconf import OmegaConf
from data.dataset import SingleFolderDataset
from loss.I2I.perceptual_loss import PerceptualLoss
import ignite.distributed as idist
CONFIG = """
loss:
perceptual:
layer_weights:
"1": 0.03125
"6": 0.0625
"11": 0.125
"20": 0.25
"29": 1
criterion: 'NL2'
style_loss: False
perceptual_loss: True
match_data:
root: "/tmp/generated/"
pipeline:
- Load
- ToTensor
- Normalize:
mean: [ 0.5, 0.5, 0.5 ]
std: [ 0.5, 0.5, 0.5 ]
not_match_data:
root: "/data/i2i/selfie2anime/trainB/"
pipeline:
- Load
- ToTensor
- Normalize:
mean: [ 0.5, 0.5, 0.5 ]
std: [ 0.5, 0.5, 0.5 ]
"""
config = OmegaConf.create(CONFIG)
dataset = SingleFolderDataset(**config.match_data)
data_loader = DataLoader(dataset, 1, False, num_workers=1)
perceptual_loss = PerceptualLoss(**config.loss.perceptual).to("cuda:0")
pls = []
for batch in data_loader:
with torch.no_grad():
batch = convert_tensor(batch, "cuda:0")
x, t = torch.chunk(batch, 2, -1)
pl, _ = perceptual_loss(x, t)
print(pl)
pls.append(pl)
torch.save(torch.stack(pls).cpu(), "verify_loss.match.pt")
dataset = SingleFolderDataset(**config.not_match_data)
data_loader = DataLoader(dataset, 4, False, num_workers=1)
pls = []
for batch in data_loader:
with torch.no_grad():
batch = convert_tensor(batch, "cuda:0")
for i, j in [(0, 1), (1, 2), (2, 3), (3, 0)]:
x, t = batch[i].unsqueeze(dim=0), batch[j].unsqueeze(dim=0)
pl, _ = perceptual_loss(x, t)
print(pl)
pls.append(pl)
torch.save(torch.stack(pls).cpu(), "verify_loss.not_match.pt")