Compare commits
24 Commits
b01016edb5
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
| 8998c30c23 | |||
| 0bec02bf6d | |||
| f7b7b78669 | |||
| 376f5caeb7 | |||
| 0019d4034c | |||
| 0927fa3de5 | |||
| 611901cbdf | |||
| a6ffab1445 | |||
| 7b05b45156 | |||
| 2de00d0245 | |||
| 74a7cfb2d8 | |||
| 436bca88b4 | |||
| 6070f08835 | |||
| 06b2abd19a | |||
| 9c08b4cd09 | |||
| 04c6366c07 | |||
| 6ea13df465 | |||
| 776fe40199 | |||
| f67bcdf161 | |||
| 16f18ab2e2 | |||
| 0f2b67e215 | |||
| acf243cb12 | |||
| fbea96f6d7 | |||
| ca55318253 |
2
.idea/misc.xml
generated
2
.idea/misc.xml
generated
@@ -1,4 +1,4 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="15d-python" project-jdk-type="Python SDK" />
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="14d-python" project-jdk-type="Python SDK" />
|
||||
</project>
|
||||
2
.idea/raycv.iml
generated
2
.idea/raycv.iml
generated
@@ -2,7 +2,7 @@
|
||||
<module type="PYTHON_MODULE" version="4">
|
||||
<component name="NewModuleRootManager">
|
||||
<content url="file://$MODULE_DIR$" />
|
||||
<orderEntry type="jdk" jdkName="15d-python" jdkType="Python SDK" />
|
||||
<orderEntry type="jdk" jdkName="14d-python" jdkType="Python SDK" />
|
||||
<orderEntry type="sourceFolder" forTests="false" />
|
||||
</component>
|
||||
<component name="TestRunnerService">
|
||||
|
||||
@@ -1,57 +1,85 @@
|
||||
name: horse2zebra-CyCleGAN
|
||||
engine: CyCleGAN
|
||||
name: huawei-cycylegan-7
|
||||
engine: CycleGAN
|
||||
result_dir: ./result
|
||||
max_pairs: 266800
|
||||
max_pairs: 1000000
|
||||
|
||||
misc:
|
||||
random_seed: 324
|
||||
|
||||
handler:
|
||||
clear_cuda_cache: False
|
||||
clear_cuda_cache: True
|
||||
set_epoch_for_dist_sampler: True
|
||||
checkpoint:
|
||||
epoch_interval: 1 # checkpoint once per `epoch_interval` epoch
|
||||
n_saved: 2
|
||||
tensorboard:
|
||||
scalar: 100 # log scalar `scalar` times per epoch
|
||||
image: 2 # log image `image` times per epoch
|
||||
image: 4 # log image `image` times per epoch
|
||||
test:
|
||||
random: True
|
||||
images: 10
|
||||
|
||||
model:
|
||||
generator:
|
||||
_type: CyCle-Generator
|
||||
_type: CycleGAN-Generator
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
out_channels: 3
|
||||
base_channels: 64
|
||||
num_blocks: 9
|
||||
padding_mode: reflect
|
||||
norm_type: IN
|
||||
use_transpose_conv: False
|
||||
pre_activation: True
|
||||
# discriminator:
|
||||
# _type: MultiScaleDiscriminator
|
||||
# _add_spectral_norm: True
|
||||
# num_scale: 2
|
||||
# down_sample_method: "bilinear"
|
||||
# discriminator_cfg:
|
||||
# _type: PatchDiscriminator
|
||||
# in_channels: 3
|
||||
# base_channels: 64
|
||||
# num_conv: 4
|
||||
# need_intermediate_feature: True
|
||||
discriminator:
|
||||
_type: PatchDiscriminator
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
base_channels: 64
|
||||
num_conv: 4
|
||||
need_intermediate_feature: False
|
||||
|
||||
|
||||
loss:
|
||||
gan:
|
||||
loss_type: lsgan
|
||||
loss_type: hinge
|
||||
weight: 1.0
|
||||
real_label_val: 1.0
|
||||
real_label_val: 1
|
||||
fake_label_val: 0.0
|
||||
cycle:
|
||||
level: 1
|
||||
weight: 10.0
|
||||
id:
|
||||
level: 1
|
||||
weight: 10.0
|
||||
mgc:
|
||||
weight: 1
|
||||
fm:
|
||||
weight: 0
|
||||
edge:
|
||||
weight: 0
|
||||
hed_pretrained_model_path: ./network-bsds500.pytorch
|
||||
|
||||
optimizers:
|
||||
generator:
|
||||
_type: Adam
|
||||
lr: 2e-4
|
||||
lr: 1e-4
|
||||
betas: [ 0.5, 0.999 ]
|
||||
weight_decay: 0.0001
|
||||
discriminator:
|
||||
_type: Adam
|
||||
lr: 2e-4
|
||||
lr: 4e-4
|
||||
betas: [ 0.5, 0.999 ]
|
||||
weight_decay: 0.0001
|
||||
|
||||
data:
|
||||
train:
|
||||
@@ -60,17 +88,28 @@ data:
|
||||
target_lr: 0
|
||||
buffer_size: 50
|
||||
dataloader:
|
||||
batch_size: 6
|
||||
batch_size: 1
|
||||
shuffle: True
|
||||
num_workers: 2
|
||||
pin_memory: True
|
||||
drop_last: True
|
||||
dataset:
|
||||
_type: GenerationUnpairedDataset
|
||||
root_a: "/data/i2i/horse2zebra/trainA"
|
||||
root_b: "/data/i2i/horse2zebra/trainB"
|
||||
root_a: "/data/face2cartoon/all_face"
|
||||
root_b: "/data/selfie2anime/trainB/"
|
||||
random_pair: True
|
||||
pipeline:
|
||||
pipeline_a:
|
||||
- Load
|
||||
- RandomCrop:
|
||||
size: [ 178, 178 ]
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- RandomHorizontalFlip
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
pipeline_b:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 286, 286 ]
|
||||
@@ -82,17 +121,38 @@ data:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
test:
|
||||
which: video_dataset
|
||||
dataloader:
|
||||
batch_size: 4
|
||||
batch_size: 1
|
||||
shuffle: False
|
||||
num_workers: 1
|
||||
pin_memory: False
|
||||
drop_last: False
|
||||
dataset:
|
||||
_type: GenerationUnpairedDataset
|
||||
root_a: "/data/i2i/horse2zebra/testA"
|
||||
root_b: "/data/i2i/horse2zebra/testB"
|
||||
random_pair: False
|
||||
root_a: "/data/face2cartoon/test/human"
|
||||
root_b: "/data/face2cartoon/test/anime"
|
||||
random_pair: True
|
||||
pipeline_a:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
pipeline_b:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
video_dataset:
|
||||
_type: SingleFolderDataset
|
||||
root: "/data/i2i/VoxCeleb2Anime/test_video_frames/"
|
||||
with_path: True
|
||||
pipeline:
|
||||
- Load
|
||||
- Resize:
|
||||
|
||||
167
configs/synthesizers/GauGAN.yml
Normal file
167
configs/synthesizers/GauGAN.yml
Normal file
@@ -0,0 +1,167 @@
|
||||
name: huawei-GauGAN-3
|
||||
engine: GauGAN
|
||||
result_dir: ./result
|
||||
max_pairs: 1000000
|
||||
|
||||
misc:
|
||||
random_seed: 324
|
||||
|
||||
handler:
|
||||
clear_cuda_cache: True
|
||||
set_epoch_for_dist_sampler: True
|
||||
checkpoint:
|
||||
epoch_interval: 1 # checkpoint once per `epoch_interval` epoch
|
||||
n_saved: 2
|
||||
tensorboard:
|
||||
scalar: 100 # log scalar `scalar` times per epoch
|
||||
image: 4 # log image `image` times per epoch
|
||||
test:
|
||||
random: True
|
||||
images: 10
|
||||
|
||||
model:
|
||||
generator:
|
||||
_type: SPADEGenerator
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
out_channels: 3
|
||||
num_blocks: 7
|
||||
use_vae: False
|
||||
num_z_dim: 256
|
||||
# discriminator:
|
||||
# _type: MultiScaleDiscriminator
|
||||
# _add_spectral_norm: True
|
||||
# num_scale: 2
|
||||
# down_sample_method: "bilinear"
|
||||
# discriminator_cfg:
|
||||
# _type: PatchDiscriminator
|
||||
# in_channels: 3
|
||||
# base_channels: 64
|
||||
# num_conv: 4
|
||||
# need_intermediate_feature: True
|
||||
discriminator:
|
||||
_type: PatchDiscriminator
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
base_channels: 64
|
||||
num_conv: 4
|
||||
need_intermediate_feature: True
|
||||
|
||||
|
||||
loss:
|
||||
gan:
|
||||
loss_type: hinge
|
||||
weight: 1.0
|
||||
real_label_val: 1
|
||||
fake_label_val: 0.0
|
||||
perceptual:
|
||||
layer_weights:
|
||||
"1": 0.03125
|
||||
"6": 0.0625
|
||||
"11": 0.125
|
||||
"20": 0.25
|
||||
"29": 1
|
||||
criterion: 'L1'
|
||||
style_loss: False
|
||||
perceptual_loss: True
|
||||
weight: 2
|
||||
mgc:
|
||||
weight: 5
|
||||
fm:
|
||||
weight: 5
|
||||
edge:
|
||||
weight: 0
|
||||
hed_pretrained_model_path: ./network-bsds500.pytorch
|
||||
|
||||
optimizers:
|
||||
generator:
|
||||
_type: Adam
|
||||
lr: 1e-4
|
||||
betas: [ 0, 0.9 ]
|
||||
weight_decay: 0.0001
|
||||
discriminator:
|
||||
_type: Adam
|
||||
lr: 4e-4
|
||||
betas: [ 0, 0.9 ]
|
||||
weight_decay: 0.0001
|
||||
|
||||
data:
|
||||
train:
|
||||
scheduler:
|
||||
start_proportion: 0.5
|
||||
target_lr: 0
|
||||
buffer_size: 50
|
||||
dataloader:
|
||||
batch_size: 1
|
||||
shuffle: True
|
||||
num_workers: 2
|
||||
pin_memory: True
|
||||
drop_last: True
|
||||
dataset:
|
||||
_type: GenerationUnpairedDataset
|
||||
root_a: "/data/face2cartoon/all_face"
|
||||
root_b: "/data/selfie2anime/trainB/"
|
||||
random_pair: True
|
||||
pipeline_a:
|
||||
- Load
|
||||
- RandomCrop:
|
||||
size: [ 178, 178 ]
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- RandomHorizontalFlip
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
pipeline_b:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 286, 286 ]
|
||||
- RandomCrop:
|
||||
size: [ 256, 256 ]
|
||||
- RandomHorizontalFlip
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
test:
|
||||
which: video_dataset
|
||||
dataloader:
|
||||
batch_size: 1
|
||||
shuffle: False
|
||||
num_workers: 1
|
||||
pin_memory: False
|
||||
drop_last: False
|
||||
dataset:
|
||||
_type: GenerationUnpairedDataset
|
||||
root_a: "/data/face2cartoon/test/human"
|
||||
root_b: "/data/face2cartoon/test/anime"
|
||||
random_pair: True
|
||||
pipeline_a:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
pipeline_b:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
video_dataset:
|
||||
_type: SingleFolderDataset
|
||||
root: "/data/i2i/VoxCeleb2Anime/test_video_frames/"
|
||||
with_path: True
|
||||
pipeline:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
@@ -19,6 +19,7 @@ handler:
|
||||
|
||||
misc:
|
||||
random_seed: 1004
|
||||
add_new_loss_epoch: -1
|
||||
|
||||
model:
|
||||
generator:
|
||||
|
||||
@@ -1,7 +1,10 @@
|
||||
name: self2anime-TSIT
|
||||
engine: TSIT
|
||||
name: huawei-TSIT-1
|
||||
engine: GauGAN
|
||||
result_dir: ./result
|
||||
max_pairs: 1500000
|
||||
max_pairs: 1000000
|
||||
|
||||
misc:
|
||||
random_seed: 324
|
||||
|
||||
handler:
|
||||
clear_cuda_cache: True
|
||||
@@ -11,36 +14,44 @@ handler:
|
||||
n_saved: 2
|
||||
tensorboard:
|
||||
scalar: 100 # log scalar `scalar` times per epoch
|
||||
image: 2 # log image `image` times per epoch
|
||||
|
||||
|
||||
misc:
|
||||
random_seed: 324
|
||||
image: 4 # log image `image` times per epoch
|
||||
test:
|
||||
random: True
|
||||
images: 10
|
||||
|
||||
model:
|
||||
generator:
|
||||
_type: TSIT-Generator
|
||||
_bn_to_sync_bn: True
|
||||
style_in_channels: 3
|
||||
content_in_channels: 3
|
||||
num_blocks: 5
|
||||
input_layer_type: "conv7x7"
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
out_channels: 3
|
||||
num_blocks: 7
|
||||
# discriminator:
|
||||
# _type: MultiScaleDiscriminator
|
||||
# _add_spectral_norm: True
|
||||
# num_scale: 2
|
||||
# down_sample_method: "bilinear"
|
||||
# discriminator_cfg:
|
||||
# _type: PatchDiscriminator
|
||||
# in_channels: 3
|
||||
# base_channels: 64
|
||||
# num_conv: 4
|
||||
# need_intermediate_feature: True
|
||||
discriminator:
|
||||
_type: MultiScaleDiscriminator
|
||||
num_scale: 2
|
||||
discriminator_cfg:
|
||||
_type: PatchDiscriminator
|
||||
in_channels: 3
|
||||
base_channels: 64
|
||||
use_spectral: True
|
||||
need_intermediate_feature: True
|
||||
_type: PatchDiscriminator
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
base_channels: 64
|
||||
num_conv: 4
|
||||
need_intermediate_feature: True
|
||||
|
||||
|
||||
loss:
|
||||
gan:
|
||||
loss_type: hinge
|
||||
real_label_val: 1.0
|
||||
fake_label_val: 0.0
|
||||
weight: 1.0
|
||||
real_label_val: 1
|
||||
fake_label_val: 0.0
|
||||
perceptual:
|
||||
layer_weights:
|
||||
"1": 0.03125
|
||||
@@ -52,25 +63,18 @@ loss:
|
||||
style_loss: False
|
||||
perceptual_loss: True
|
||||
weight: 1
|
||||
style:
|
||||
layer_weights:
|
||||
"1": 0.03125
|
||||
"6": 0.0625
|
||||
"11": 0.125
|
||||
"20": 0.25
|
||||
"29": 1
|
||||
criterion: 'L2'
|
||||
style_loss: True
|
||||
perceptual_loss: False
|
||||
weight: 0
|
||||
mgc:
|
||||
weight: 5
|
||||
fm:
|
||||
level: 1
|
||||
weight: 1
|
||||
edge:
|
||||
weight: 0
|
||||
hed_pretrained_model_path: ./network-bsds500.pytorch
|
||||
|
||||
optimizers:
|
||||
generator:
|
||||
_type: Adam
|
||||
lr: 0.0001
|
||||
lr: 1e-4
|
||||
betas: [ 0, 0.9 ]
|
||||
weight_decay: 0.0001
|
||||
discriminator:
|
||||
@@ -84,7 +88,7 @@ data:
|
||||
scheduler:
|
||||
start_proportion: 0.5
|
||||
target_lr: 0
|
||||
buffer_size: 50
|
||||
buffer_size: 0
|
||||
dataloader:
|
||||
batch_size: 1
|
||||
shuffle: True
|
||||
@@ -92,18 +96,28 @@ data:
|
||||
pin_memory: True
|
||||
drop_last: True
|
||||
dataset:
|
||||
_type: GenerationUnpairedDatasetWithEdge
|
||||
root_a: "/data/i2i/VoxCeleb2Anime/trainA"
|
||||
root_b: "/data/i2i/VoxCeleb2Anime/trainB"
|
||||
edges_path: "/data/i2i/VoxCeleb2Anime/edges"
|
||||
landmarks_path: "/data/i2i/VoxCeleb2Anime/landmarks"
|
||||
edge_type: "landmark_hed"
|
||||
size: [ 128, 128 ]
|
||||
_type: GenerationUnpairedDataset
|
||||
root_a: "/data/face2cartoon/all_face"
|
||||
root_b: "/data/selfie2anime/trainB/"
|
||||
random_pair: True
|
||||
pipeline:
|
||||
pipeline_a:
|
||||
- Load
|
||||
- RandomCrop:
|
||||
size: [ 178, 178 ]
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- RandomHorizontalFlip
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
pipeline_b:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 128, 128 ]
|
||||
size: [ 286, 286 ]
|
||||
- RandomCrop:
|
||||
size: [ 256, 256 ]
|
||||
- RandomHorizontalFlip
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
@@ -111,21 +125,28 @@ data:
|
||||
test:
|
||||
which: video_dataset
|
||||
dataloader:
|
||||
batch_size: 8
|
||||
batch_size: 1
|
||||
shuffle: False
|
||||
num_workers: 1
|
||||
pin_memory: False
|
||||
drop_last: False
|
||||
dataset:
|
||||
_type: GenerationUnpairedDataset
|
||||
root_a: "/data/i2i/VoxCeleb2Anime/testA"
|
||||
root_b: "/data/i2i/VoxCeleb2Anime/testB"
|
||||
with_path: True
|
||||
random_pair: False
|
||||
pipeline:
|
||||
root_a: "/data/face2cartoon/test/human"
|
||||
root_b: "/data/face2cartoon/test/anime"
|
||||
random_pair: True
|
||||
pipeline_a:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 128, 128 ]
|
||||
size: [ 256, 256 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
pipeline_b:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 256, 256 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
|
||||
@@ -14,11 +14,15 @@ handler:
|
||||
n_saved: 2
|
||||
tensorboard:
|
||||
scalar: 100 # log scalar `scalar` times per epoch
|
||||
image: 2 # log image `image` times per epoch
|
||||
image: 4 # log image `image` times per epoch
|
||||
test:
|
||||
random: True
|
||||
images: 10
|
||||
|
||||
model:
|
||||
generator:
|
||||
_type: UGATIT-Generator
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
out_channels: 3
|
||||
base_channels: 64
|
||||
@@ -27,11 +31,13 @@ model:
|
||||
light: True
|
||||
local_discriminator:
|
||||
_type: UGATIT-Discriminator
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
base_channels: 64
|
||||
num_blocks: 5
|
||||
global_discriminator:
|
||||
_type: UGATIT-Discriminator
|
||||
_add_spectral_norm: True
|
||||
in_channels: 3
|
||||
base_channels: 64
|
||||
num_blocks: 7
|
||||
@@ -50,6 +56,8 @@ loss:
|
||||
weight: 10.0
|
||||
cam:
|
||||
weight: 1000
|
||||
mgc:
|
||||
weight: 0
|
||||
|
||||
optimizers:
|
||||
generator:
|
||||
@@ -70,7 +78,7 @@ data:
|
||||
target_lr: 0
|
||||
buffer_size: 50
|
||||
dataloader:
|
||||
batch_size: 24
|
||||
batch_size: 1
|
||||
shuffle: True
|
||||
num_workers: 2
|
||||
pin_memory: True
|
||||
@@ -94,7 +102,7 @@ data:
|
||||
test:
|
||||
which: video_dataset
|
||||
dataloader:
|
||||
batch_size: 8
|
||||
batch_size: 1
|
||||
shuffle: False
|
||||
num_workers: 1
|
||||
pin_memory: False
|
||||
|
||||
171
configs/synthesizers/talking_anime.yml
Normal file
171
configs/synthesizers/talking_anime.yml
Normal file
@@ -0,0 +1,171 @@
|
||||
name: talking_anime
|
||||
engine: talking_anime
|
||||
result_dir: ./result
|
||||
max_pairs: 1000000
|
||||
|
||||
handler:
|
||||
clear_cuda_cache: True
|
||||
set_epoch_for_dist_sampler: True
|
||||
checkpoint:
|
||||
epoch_interval: 1 # checkpoint once per `epoch_interval` epoch
|
||||
n_saved: 2
|
||||
tensorboard:
|
||||
scalar: 100 # log scalar `scalar` times per epoch
|
||||
image: 100 # log image `image` times per epoch
|
||||
test:
|
||||
random: True
|
||||
images: 10
|
||||
|
||||
misc:
|
||||
random_seed: 1004
|
||||
|
||||
loss:
|
||||
gan:
|
||||
loss_type: hinge
|
||||
real_label_val: 1.0
|
||||
fake_label_val: 0.0
|
||||
weight: 1.0
|
||||
fm:
|
||||
level: 1
|
||||
weight: 1
|
||||
style:
|
||||
layer_weights:
|
||||
"3": 1
|
||||
criterion: 'L1'
|
||||
style_loss: True
|
||||
perceptual_loss: False
|
||||
weight: 10
|
||||
perceptual:
|
||||
layer_weights:
|
||||
"1": 0.03125
|
||||
"6": 0.0625
|
||||
"11": 0.125
|
||||
"20": 0.25
|
||||
"29": 1
|
||||
criterion: 'L1'
|
||||
style_loss: False
|
||||
perceptual_loss: True
|
||||
weight: 0
|
||||
context:
|
||||
layer_weights:
|
||||
#"13": 1
|
||||
"22": 1
|
||||
weight: 5
|
||||
recon:
|
||||
level: 1
|
||||
weight: 10
|
||||
edge:
|
||||
weight: 5
|
||||
hed_pretrained_model_path: ./network-bsds500.pytorch
|
||||
|
||||
model:
|
||||
face_generator:
|
||||
_type: TAFG-SingleGenerator
|
||||
_bn_to_sync_bn: False
|
||||
style_in_channels: 3
|
||||
content_in_channels: 1
|
||||
use_spectral_norm: True
|
||||
style_encoder_type: VGG19StyleEncoder
|
||||
num_style_conv: 4
|
||||
style_dim: 512
|
||||
num_adain_blocks: 8
|
||||
num_res_blocks: 8
|
||||
anime_generator:
|
||||
_type: TAFG-ResGenerator
|
||||
_bn_to_sync_bn: False
|
||||
in_channels: 6
|
||||
use_spectral_norm: True
|
||||
num_res_blocks: 8
|
||||
|
||||
discriminator:
|
||||
_type: MultiScaleDiscriminator
|
||||
num_scale: 2
|
||||
discriminator_cfg:
|
||||
_type: PatchDiscriminator
|
||||
in_channels: 3
|
||||
base_channels: 64
|
||||
use_spectral: True
|
||||
need_intermediate_feature: True
|
||||
|
||||
optimizers:
|
||||
generator:
|
||||
_type: Adam
|
||||
lr: 0.0001
|
||||
betas: [ 0, 0.9 ]
|
||||
weight_decay: 0.0001
|
||||
discriminator:
|
||||
_type: Adam
|
||||
lr: 4e-4
|
||||
betas: [ 0, 0.9 ]
|
||||
weight_decay: 0.0001
|
||||
|
||||
data:
|
||||
train:
|
||||
scheduler:
|
||||
start_proportion: 0.5
|
||||
target_lr: 0
|
||||
dataloader:
|
||||
batch_size: 8
|
||||
shuffle: True
|
||||
num_workers: 1
|
||||
pin_memory: True
|
||||
drop_last: True
|
||||
dataset:
|
||||
_type: PoseFacesWithSingleAnime
|
||||
root_face: "/data/i2i/VoxCeleb2Anime/trainA"
|
||||
root_anime: "/data/i2i/VoxCeleb2Anime/trainB"
|
||||
landmark_path: "/data/i2i/VoxCeleb2Anime/landmarks"
|
||||
num_face: 2
|
||||
img_size: [ 128, 128 ]
|
||||
with_order: False
|
||||
face_pipeline:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 128, 128 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
anime_pipeline:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 144, 144 ]
|
||||
- RandomCrop:
|
||||
size: [ 128, 128 ]
|
||||
- RandomHorizontalFlip
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
test:
|
||||
which: dataset
|
||||
dataloader:
|
||||
batch_size: 1
|
||||
shuffle: False
|
||||
num_workers: 1
|
||||
pin_memory: False
|
||||
drop_last: False
|
||||
dataset:
|
||||
_type: PoseFacesWithSingleAnime
|
||||
root_face: "/data/i2i/VoxCeleb2Anime/testA"
|
||||
root_anime: "/data/i2i/VoxCeleb2Anime/testB"
|
||||
landmark_path: "/data/i2i/VoxCeleb2Anime/landmarks"
|
||||
num_face: 2
|
||||
img_size: [ 128, 128 ]
|
||||
with_order: False
|
||||
face_pipeline:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 128, 128 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
anime_pipeline:
|
||||
- Load
|
||||
- Resize:
|
||||
size: [ 128, 128 ]
|
||||
- ToTensor
|
||||
- Normalize:
|
||||
mean: [ 0.5, 0.5, 0.5 ]
|
||||
std: [ 0.5, 0.5, 0.5 ]
|
||||
239
data/dataset.py
239
data/dataset.py
@@ -1,239 +0,0 @@
|
||||
import os
|
||||
import pickle
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
import lmdb
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision.datasets import ImageFolder
|
||||
from torchvision.datasets.folder import has_file_allowed_extension, IMG_EXTENSIONS
|
||||
from torchvision.transforms import functional as F
|
||||
from tqdm import tqdm
|
||||
|
||||
from .registry import DATASET
|
||||
from .transform import transform_pipeline
|
||||
from .util import dlib_landmark
|
||||
|
||||
|
||||
def default_transform_way(transform, sample):
|
||||
return [transform(sample[0]), *sample[1:]]
|
||||
|
||||
|
||||
class LMDBDataset(Dataset):
|
||||
def __init__(self, lmdb_path, pipeline=None, transform_way=default_transform_way, map_size=2 ** 40, readonly=True,
|
||||
**lmdb_kwargs):
|
||||
self.path = lmdb_path
|
||||
self.db = lmdb.open(lmdb_path, subdir=os.path.isdir(lmdb_path), map_size=map_size, readonly=readonly,
|
||||
lock=False, **lmdb_kwargs)
|
||||
|
||||
with self.db.begin(write=False) as txn:
|
||||
self._len = pickle.loads(txn.get(b"$$len$$"))
|
||||
self.done_pipeline = pickle.loads(txn.get(b"$$done_pipeline$$"))
|
||||
if pipeline is None:
|
||||
self.not_done_pipeline = []
|
||||
else:
|
||||
self.not_done_pipeline = self._remain_pipeline(pipeline)
|
||||
self.transform = transform_pipeline(self.not_done_pipeline)
|
||||
self.transform_way = transform_way
|
||||
essential_attr = pickle.loads(txn.get(b"$$essential_attr$$"))
|
||||
for ea in essential_attr:
|
||||
setattr(self, ea, pickle.loads(txn.get(f"${ea}$".encode(encoding="utf-8"))))
|
||||
|
||||
def _remain_pipeline(self, pipeline):
|
||||
for i, dp in enumerate(self.done_pipeline):
|
||||
if pipeline[i] != dp:
|
||||
raise ValueError(
|
||||
f"pipeline {self.done_pipeline} saved in this lmdb database is not match with pipeline:{pipeline}")
|
||||
return pipeline[len(self.done_pipeline):]
|
||||
|
||||
def __repr__(self):
|
||||
return f"LMDBDataset: {self.path}\nlength: {len(self)}\n{self.transform}"
|
||||
|
||||
def __len__(self):
|
||||
return self._len
|
||||
|
||||
def __getitem__(self, idx):
|
||||
with self.db.begin(write=False) as txn:
|
||||
sample = pickle.loads(txn.get("{}".format(idx).encode()))
|
||||
sample = self.transform_way(self.transform, sample)
|
||||
return sample
|
||||
|
||||
@staticmethod
|
||||
def lmdbify(dataset, done_pipeline, lmdb_path):
|
||||
env = lmdb.open(lmdb_path, map_size=2 ** 40, subdir=os.path.isdir(lmdb_path))
|
||||
with env.begin(write=True) as txn:
|
||||
for i in tqdm(range(len(dataset)), ncols=0):
|
||||
txn.put("{}".format(i).encode(), pickle.dumps(dataset[i]))
|
||||
txn.put(b"$$len$$", pickle.dumps(len(dataset)))
|
||||
txn.put(b"$$done_pipeline$$", pickle.dumps(done_pipeline))
|
||||
essential_attr = getattr(dataset, "essential_attr", list())
|
||||
txn.put(b"$$essential_attr$$", pickle.dumps(essential_attr))
|
||||
for ea in essential_attr:
|
||||
txn.put(f"${ea}$".encode(encoding="utf-8"), pickle.dumps(getattr(dataset, ea)))
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class ImprovedImageFolder(ImageFolder):
|
||||
def __init__(self, root, pipeline):
|
||||
super().__init__(root, transform_pipeline(pipeline), loader=lambda x: x)
|
||||
self.classes_list = defaultdict(list)
|
||||
self.essential_attr = ["classes_list"]
|
||||
for i in range(len(self)):
|
||||
self.classes_list[self.samples[i][-1]].append(i)
|
||||
assert len(self.classes_list) == len(self.classes)
|
||||
|
||||
|
||||
class EpisodicDataset(Dataset):
|
||||
def __init__(self, origin_dataset, num_class, num_query, num_support, num_episodes):
|
||||
self.origin = origin_dataset
|
||||
self.num_class = num_class
|
||||
assert self.num_class < len(self.origin.classes_list)
|
||||
self.num_query = num_query # K
|
||||
self.num_support = num_support # K
|
||||
self.num_episodes = num_episodes
|
||||
|
||||
def _fetch_list_data(self, id_list):
|
||||
return [self.origin[i][0] for i in id_list]
|
||||
|
||||
def __len__(self):
|
||||
return self.num_episodes
|
||||
|
||||
def __getitem__(self, _):
|
||||
random_classes = torch.randperm(len(self.origin.classes_list))[:self.num_class].tolist()
|
||||
support_set = []
|
||||
query_set = []
|
||||
target_set = []
|
||||
for tag, c in enumerate(random_classes):
|
||||
image_list = self.origin.classes_list[c]
|
||||
|
||||
if len(image_list) >= self.num_query + self.num_support:
|
||||
# have enough images belong to this class
|
||||
idx_list = torch.randperm(len(image_list))[:self.num_query + self.num_support].tolist()
|
||||
else:
|
||||
idx_list = torch.randint(high=len(image_list), size=(self.num_query + self.num_support,)).tolist()
|
||||
|
||||
support = self._fetch_list_data(map(image_list.__getitem__, idx_list[:self.num_support]))
|
||||
query = self._fetch_list_data(map(image_list.__getitem__, idx_list[self.num_support:]))
|
||||
support_set.extend(support)
|
||||
query_set.extend(query)
|
||||
target_set.extend([tag] * self.num_query)
|
||||
return {
|
||||
"support": torch.stack(support_set),
|
||||
"query": torch.stack(query_set),
|
||||
"target": torch.tensor(target_set)
|
||||
}
|
||||
|
||||
def __repr__(self):
|
||||
return f"<EpisodicDataset NE{self.num_episodes} NC{self.num_class} NS{self.num_support} NQ{self.num_query}>"
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class SingleFolderDataset(Dataset):
|
||||
def __init__(self, root, pipeline, with_path=False):
|
||||
assert os.path.isdir(root)
|
||||
self.root = root
|
||||
samples = []
|
||||
for r, _, fns in sorted(os.walk(self.root, followlinks=True)):
|
||||
for fn in sorted(fns):
|
||||
path = os.path.join(r, fn)
|
||||
if has_file_allowed_extension(path, IMG_EXTENSIONS):
|
||||
samples.append(path)
|
||||
self.samples = samples
|
||||
self.pipeline = transform_pipeline(pipeline)
|
||||
self.with_path = with_path
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if not self.with_path:
|
||||
return self.pipeline(self.samples[idx])
|
||||
else:
|
||||
return self.pipeline(self.samples[idx]), self.samples[idx]
|
||||
|
||||
def __repr__(self):
|
||||
return f"<SingleFolderDataset root={self.root} len={len(self)}>"
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class GenerationUnpairedDataset(Dataset):
|
||||
def __init__(self, root_a, root_b, random_pair, pipeline, with_path=False):
|
||||
self.A = SingleFolderDataset(root_a, pipeline, with_path)
|
||||
self.B = SingleFolderDataset(root_b, pipeline, with_path)
|
||||
self.random_pair = random_pair
|
||||
|
||||
def __getitem__(self, idx):
|
||||
a_idx = idx % len(self.A)
|
||||
b_idx = idx % len(self.B) if not self.random_pair else torch.randint(len(self.B), (1,)).item()
|
||||
return dict(a=self.A[a_idx], b=self.B[b_idx])
|
||||
|
||||
def __len__(self):
|
||||
return max(len(self.A), len(self.B))
|
||||
|
||||
def __repr__(self):
|
||||
return f"<GenerationUnpairedDataset:\n\tA: {self.A}\n\tB: {self.B}>\nPipeline:\n{self.A.pipeline}"
|
||||
|
||||
|
||||
def normalize_tensor(tensor):
|
||||
tensor = tensor.float()
|
||||
tensor -= tensor.min()
|
||||
tensor /= tensor.max()
|
||||
return tensor
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class GenerationUnpairedDatasetWithEdge(Dataset):
|
||||
def __init__(self, root_a, root_b, random_pair, pipeline, edge_type, edges_path, landmarks_path, size=(256, 256),
|
||||
with_path=False):
|
||||
assert edge_type in ["hed", "canny", "landmark_hed", "landmark_canny"]
|
||||
self.edge_type = edge_type
|
||||
self.size = size
|
||||
self.edges_path = Path(edges_path)
|
||||
self.landmarks_path = Path(landmarks_path)
|
||||
assert self.edges_path.exists()
|
||||
assert self.landmarks_path.exists()
|
||||
self.A = SingleFolderDataset(root_a, pipeline, with_path=True)
|
||||
self.B = SingleFolderDataset(root_b, pipeline, with_path=True)
|
||||
self.random_pair = random_pair
|
||||
self.with_path = with_path
|
||||
|
||||
def get_edge(self, origin_path):
|
||||
op = Path(origin_path)
|
||||
if self.edge_type.startswith("landmark_"):
|
||||
edge_type = self.edge_type.lstrip("landmark_")
|
||||
use_landmark = op.parent.name.endswith("A")
|
||||
else:
|
||||
edge_type = self.edge_type
|
||||
use_landmark = False
|
||||
edge_path = self.edges_path / f"{op.parent.name}/{op.stem}.{edge_type}.png"
|
||||
origin_edge = F.to_tensor(Image.open(edge_path).resize(self.size, Image.BILINEAR))
|
||||
if not use_landmark:
|
||||
return origin_edge
|
||||
else:
|
||||
landmark_path = self.landmarks_path / f"{op.parent.name}/{op.stem}.txt"
|
||||
key_points, part_labels, part_edge = dlib_landmark.read_keypoints(landmark_path, size=self.size)
|
||||
|
||||
dist_tensor = normalize_tensor(torch.from_numpy(dlib_landmark.dist_tensor(key_points, size=self.size)))
|
||||
part_labels = normalize_tensor(torch.from_numpy(part_labels))
|
||||
part_edge = torch.from_numpy(part_edge).unsqueeze(0).float()
|
||||
# edges = origin_edge * (part_labels.sum(0) == 0) # remove edges within face
|
||||
# edges = part_edge + edges
|
||||
return torch.cat([origin_edge, part_edge, dist_tensor, part_labels])
|
||||
|
||||
def __getitem__(self, idx):
|
||||
a_idx = idx % len(self.A)
|
||||
b_idx = idx % len(self.B) if not self.random_pair else torch.randint(len(self.B), (1,)).item()
|
||||
output = dict(a={}, b={})
|
||||
output["a"]["img"], output["a"]["path"] = self.A[a_idx]
|
||||
output["b"]["img"], output["b"]["path"] = self.B[b_idx]
|
||||
for p in "ab":
|
||||
output[p]["edge"] = self.get_edge(output[p]["path"])
|
||||
return output
|
||||
|
||||
def __len__(self):
|
||||
return max(len(self.A), len(self.B))
|
||||
|
||||
def __repr__(self):
|
||||
return f"<GenerationUnpairedDatasetWithEdge:\n\tA: {self.A}\n\tB: {self.B}>\nPipeline:\n{self.A.pipeline}"
|
||||
3
data/dataset/__init__.py
Normal file
3
data/dataset/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
from util.misc import import_submodules
|
||||
|
||||
__all__ = import_submodules(__name__).keys()
|
||||
63
data/dataset/few-shot.py
Normal file
63
data/dataset/few-shot.py
Normal file
@@ -0,0 +1,63 @@
|
||||
from collections import defaultdict
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision.datasets import ImageFolder
|
||||
|
||||
from data.registry import DATASET
|
||||
from data.transform import transform_pipeline
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class ImprovedImageFolder(ImageFolder):
|
||||
def __init__(self, root, pipeline):
|
||||
super().__init__(root, transform_pipeline(pipeline), loader=lambda x: x)
|
||||
self.classes_list = defaultdict(list)
|
||||
self.essential_attr = ["classes_list"]
|
||||
for i in range(len(self)):
|
||||
self.classes_list[self.samples[i][-1]].append(i)
|
||||
assert len(self.classes_list) == len(self.classes)
|
||||
|
||||
|
||||
class EpisodicDataset(Dataset):
|
||||
def __init__(self, origin_dataset, num_class, num_query, num_support, num_episodes):
|
||||
self.origin = origin_dataset
|
||||
self.num_class = num_class
|
||||
assert self.num_class < len(self.origin.classes_list)
|
||||
self.num_query = num_query # K
|
||||
self.num_support = num_support # K
|
||||
self.num_episodes = num_episodes
|
||||
|
||||
def _fetch_list_data(self, id_list):
|
||||
return [self.origin[i][0] for i in id_list]
|
||||
|
||||
def __len__(self):
|
||||
return self.num_episodes
|
||||
|
||||
def __getitem__(self, _):
|
||||
random_classes = torch.randperm(len(self.origin.classes_list))[:self.num_class].tolist()
|
||||
support_set = []
|
||||
query_set = []
|
||||
target_set = []
|
||||
for tag, c in enumerate(random_classes):
|
||||
image_list = self.origin.classes_list[c]
|
||||
|
||||
if len(image_list) >= self.num_query + self.num_support:
|
||||
# have enough images belong to this class
|
||||
idx_list = torch.randperm(len(image_list))[:self.num_query + self.num_support].tolist()
|
||||
else:
|
||||
idx_list = torch.randint(high=len(image_list), size=(self.num_query + self.num_support,)).tolist()
|
||||
|
||||
support = self._fetch_list_data(map(image_list.__getitem__, idx_list[:self.num_support]))
|
||||
query = self._fetch_list_data(map(image_list.__getitem__, idx_list[self.num_support:]))
|
||||
support_set.extend(support)
|
||||
query_set.extend(query)
|
||||
target_set.extend([tag] * self.num_query)
|
||||
return {
|
||||
"support": torch.stack(support_set),
|
||||
"query": torch.stack(query_set),
|
||||
"target": torch.tensor(target_set)
|
||||
}
|
||||
|
||||
def __repr__(self):
|
||||
return f"<EpisodicDataset NE{self.num_episodes} NC{self.num_class} NS{self.num_support} NQ{self.num_query}>"
|
||||
62
data/dataset/image_translation.py
Normal file
62
data/dataset/image_translation.py
Normal file
@@ -0,0 +1,62 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision.datasets.folder import has_file_allowed_extension, IMG_EXTENSIONS
|
||||
|
||||
from data.registry import DATASET
|
||||
from data.transform import transform_pipeline
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class SingleFolderDataset(Dataset):
|
||||
def __init__(self, root, pipeline, with_path=False):
|
||||
assert os.path.isdir(root)
|
||||
self.root = root
|
||||
samples = []
|
||||
for r, _, fns in sorted(os.walk(self.root, followlinks=True)):
|
||||
for fn in sorted(fns):
|
||||
path = os.path.join(r, fn)
|
||||
if has_file_allowed_extension(path, IMG_EXTENSIONS):
|
||||
samples.append(path)
|
||||
self.samples = samples
|
||||
self.pipeline = transform_pipeline(pipeline)
|
||||
self.with_path = with_path
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
output = dict(img=self.pipeline(self.samples[idx]))
|
||||
if self.with_path:
|
||||
output["path"] = self.samples[idx]
|
||||
return output
|
||||
|
||||
def __repr__(self):
|
||||
return f"<SingleFolderDataset root={self.root} len={len(self)} with_path={self.with_path}>"
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class GenerationUnpairedDataset(Dataset):
|
||||
def __init__(self, root_a, root_b, random_pair, pipeline_a, pipeline_b, with_path=False):
|
||||
self.A = SingleFolderDataset(root_a, pipeline_a, with_path)
|
||||
self.B = SingleFolderDataset(root_b, pipeline_b, with_path)
|
||||
self.with_path = with_path
|
||||
self.random_pair = random_pair
|
||||
|
||||
def __getitem__(self, idx):
|
||||
a_idx = idx % len(self.A)
|
||||
b_idx = idx % len(self.B) if not self.random_pair else torch.randint(len(self.B), (1,)).item()
|
||||
output_a = self.A[a_idx]
|
||||
output_b = self.B[b_idx]
|
||||
output = dict(a=output_a["img"], b=output_b["img"])
|
||||
if self.with_path:
|
||||
output["a_path"] = output_a["path"]
|
||||
output["b_path"] = output_b["path"]
|
||||
return output
|
||||
|
||||
def __len__(self):
|
||||
return max(len(self.A), len(self.B))
|
||||
|
||||
def __repr__(self):
|
||||
return f"<GenerationUnpairedDataset:\n\tA: {self.A}\n\tB: {self.B}>\nPipeline:\n{self.A.pipeline}"
|
||||
65
data/dataset/lmdb.py
Normal file
65
data/dataset/lmdb.py
Normal file
@@ -0,0 +1,65 @@
|
||||
import os
|
||||
import pickle
|
||||
|
||||
import lmdb
|
||||
from torch.utils.data import Dataset
|
||||
from tqdm import tqdm
|
||||
|
||||
from data.transform import transform_pipeline
|
||||
|
||||
|
||||
def default_transform_way(transform, sample):
|
||||
return [transform(sample[0]), *sample[1:]]
|
||||
|
||||
|
||||
class LMDBDataset(Dataset):
|
||||
def __init__(self, lmdb_path, pipeline=None, transform_way=default_transform_way, map_size=2 ** 40, readonly=True,
|
||||
**lmdb_kwargs):
|
||||
self.path = lmdb_path
|
||||
self.db = lmdb.open(lmdb_path, subdir=os.path.isdir(lmdb_path), map_size=map_size, readonly=readonly,
|
||||
lock=False, **lmdb_kwargs)
|
||||
|
||||
with self.db.begin(write=False) as txn:
|
||||
self._len = pickle.loads(txn.get(b"$$len$$"))
|
||||
self.done_pipeline = pickle.loads(txn.get(b"$$done_pipeline$$"))
|
||||
if pipeline is None:
|
||||
self.not_done_pipeline = []
|
||||
else:
|
||||
self.not_done_pipeline = self._remain_pipeline(pipeline)
|
||||
self.transform = transform_pipeline(self.not_done_pipeline)
|
||||
self.transform_way = transform_way
|
||||
essential_attr = pickle.loads(txn.get(b"$$essential_attr$$"))
|
||||
for ea in essential_attr:
|
||||
setattr(self, ea, pickle.loads(txn.get(f"${ea}$".encode(encoding="utf-8"))))
|
||||
|
||||
def _remain_pipeline(self, pipeline):
|
||||
for i, dp in enumerate(self.done_pipeline):
|
||||
if pipeline[i] != dp:
|
||||
raise ValueError(
|
||||
f"pipeline {self.done_pipeline} saved in this lmdb database is not match with pipeline:{pipeline}")
|
||||
return pipeline[len(self.done_pipeline):]
|
||||
|
||||
def __repr__(self):
|
||||
return f"LMDBDataset: {self.path}\nlength: {len(self)}\n{self.transform}"
|
||||
|
||||
def __len__(self):
|
||||
return self._len
|
||||
|
||||
def __getitem__(self, idx):
|
||||
with self.db.begin(write=False) as txn:
|
||||
sample = pickle.loads(txn.get("{}".format(idx).encode()))
|
||||
sample = self.transform_way(self.transform, sample)
|
||||
return sample
|
||||
|
||||
@staticmethod
|
||||
def lmdbify(dataset, done_pipeline, lmdb_path):
|
||||
env = lmdb.open(lmdb_path, map_size=2 ** 40, subdir=os.path.isdir(lmdb_path))
|
||||
with env.begin(write=True) as txn:
|
||||
for i in tqdm(range(len(dataset)), ncols=0):
|
||||
txn.put("{}".format(i).encode(), pickle.dumps(dataset[i]))
|
||||
txn.put(b"$$len$$", pickle.dumps(len(dataset)))
|
||||
txn.put(b"$$done_pipeline$$", pickle.dumps(done_pipeline))
|
||||
essential_attr = getattr(dataset, "essential_attr", list())
|
||||
txn.put(b"$$essential_attr$$", pickle.dumps(essential_attr))
|
||||
for ea in essential_attr:
|
||||
txn.put(f"${ea}$".encode(encoding="utf-8"), pickle.dumps(getattr(dataset, ea)))
|
||||
122
data/dataset/pose_transfer.py
Normal file
122
data/dataset/pose_transfer.py
Normal file
@@ -0,0 +1,122 @@
|
||||
from collections import defaultdict
|
||||
from itertools import permutations, combinations
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision.transforms import functional as F
|
||||
|
||||
from data.registry import DATASET
|
||||
from data.transform import transform_pipeline
|
||||
from data.util import dlib_landmark
|
||||
|
||||
|
||||
def normalize_tensor(tensor):
|
||||
tensor = tensor.float()
|
||||
tensor -= tensor.min()
|
||||
tensor /= tensor.max()
|
||||
return tensor
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class GenerationUnpairedDatasetWithEdge(Dataset):
|
||||
def __init__(self, root_a, root_b, random_pair, pipeline, edge_type, edges_path, landmarks_path, size=(256, 256),
|
||||
with_path=False):
|
||||
assert edge_type in ["hed", "canny", "landmark_hed", "landmark_canny"]
|
||||
self.edge_type = edge_type
|
||||
self.size = size
|
||||
self.edges_path = Path(edges_path)
|
||||
self.landmarks_path = Path(landmarks_path)
|
||||
assert self.edges_path.exists()
|
||||
assert self.landmarks_path.exists()
|
||||
self.A = SingleFolderDataset(root_a, pipeline, with_path=True)
|
||||
self.B = SingleFolderDataset(root_b, pipeline, with_path=True)
|
||||
self.random_pair = random_pair
|
||||
self.with_path = with_path
|
||||
|
||||
def get_edge(self, origin_path):
|
||||
op = Path(origin_path)
|
||||
if self.edge_type.startswith("landmark_"):
|
||||
edge_type = self.edge_type.lstrip("landmark_")
|
||||
use_landmark = op.parent.name.endswith("A")
|
||||
else:
|
||||
edge_type = self.edge_type
|
||||
use_landmark = False
|
||||
edge_path = self.edges_path / f"{op.parent.name}/{op.stem}.{edge_type}.png"
|
||||
origin_edge = F.to_tensor(Image.open(edge_path).resize(self.size, Image.BILINEAR))
|
||||
if not use_landmark:
|
||||
return origin_edge
|
||||
else:
|
||||
landmark_path = self.landmarks_path / f"{op.parent.name}/{op.stem}.txt"
|
||||
key_points, part_labels, part_edge = dlib_landmark.read_keypoints(landmark_path, size=self.size)
|
||||
|
||||
dist_tensor = normalize_tensor(torch.from_numpy(dlib_landmark.dist_tensor(key_points, size=self.size)))
|
||||
part_labels = normalize_tensor(torch.from_numpy(part_labels))
|
||||
part_edge = torch.from_numpy(part_edge).unsqueeze(0).float()
|
||||
# edges = origin_edge * (part_labels.sum(0) == 0) # remove edges within face
|
||||
# edges = part_edge + edges
|
||||
return torch.cat([origin_edge, part_edge, dist_tensor, part_labels])
|
||||
|
||||
def __getitem__(self, idx):
|
||||
a_idx = idx % len(self.A)
|
||||
b_idx = idx % len(self.B) if not self.random_pair else torch.randint(len(self.B), (1,)).item()
|
||||
output = dict(a={}, b={})
|
||||
output["a"]["img"], output["a"]["path"] = self.A[a_idx]
|
||||
output["b"]["img"], output["b"]["path"] = self.B[b_idx]
|
||||
for p in "ab":
|
||||
output[p]["edge"] = self.get_edge(output[p]["path"])
|
||||
return output
|
||||
|
||||
def __len__(self):
|
||||
return max(len(self.A), len(self.B))
|
||||
|
||||
def __repr__(self):
|
||||
return f"<GenerationUnpairedDatasetWithEdge:\n\tA: {self.A}\n\tB: {self.B}>\nPipeline:\n{self.A.pipeline}"
|
||||
|
||||
|
||||
@DATASET.register_module()
|
||||
class PoseFacesWithSingleAnime(Dataset):
|
||||
def __init__(self, root_face, root_anime, landmark_path, num_face, face_pipeline, anime_pipeline, img_size,
|
||||
with_order=True):
|
||||
self.num_face = num_face
|
||||
self.landmark_path = Path(landmark_path)
|
||||
self.with_order = with_order
|
||||
self.root_face = Path(root_face)
|
||||
self.root_anime = Path(root_anime)
|
||||
self.img_size = img_size
|
||||
self.face_samples = self.iter_folders()
|
||||
self.face_pipeline = transform_pipeline(face_pipeline)
|
||||
self.B = SingleFolderDataset(root_anime, anime_pipeline, with_path=True)
|
||||
|
||||
def iter_folders(self):
|
||||
pics_per_person = defaultdict(list)
|
||||
for p in self.root_face.glob("*.jpg"):
|
||||
pics_per_person[p.stem[:7]].append(p.stem)
|
||||
data = []
|
||||
for p in pics_per_person:
|
||||
if len(pics_per_person[p]) >= self.num_face:
|
||||
if self.with_order:
|
||||
data.extend(list(combinations(pics_per_person[p], self.num_face)))
|
||||
else:
|
||||
data.extend(list(permutations(pics_per_person[p], self.num_face)))
|
||||
return data
|
||||
|
||||
def read_pose(self, pose_txt):
|
||||
key_points, part_labels, part_edge = dlib_landmark.read_keypoints(pose_txt, size=self.img_size)
|
||||
dist_tensor = normalize_tensor(torch.from_numpy(dlib_landmark.dist_tensor(key_points, size=self.img_size)))
|
||||
part_labels = normalize_tensor(torch.from_numpy(part_labels))
|
||||
part_edge = torch.from_numpy(part_edge).unsqueeze(0).float()
|
||||
return torch.cat([part_labels, part_edge, dist_tensor])
|
||||
|
||||
def __len__(self):
|
||||
return len(self.face_samples)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
output = dict()
|
||||
output["anime_img"], output["anime_path"] = self.B[torch.randint(len(self.B), (1,)).item()]
|
||||
for i, f in enumerate(self.face_samples[idx]):
|
||||
output[f"face_{i}"] = self.face_pipeline(self.root_face / f"{f}.jpg")
|
||||
output[f"pose_{i}"] = self.read_pose(self.landmark_path / self.root_face.name / f"{f}.txt")
|
||||
output[f"stem_{i}"] = f
|
||||
return output
|
||||
@@ -28,7 +28,7 @@ class Load:
|
||||
|
||||
|
||||
def transform_pipeline(pipeline_description):
|
||||
if len(pipeline_description) == 0:
|
||||
if pipeline_description is None or len(pipeline_description) == 0:
|
||||
return lambda x: x
|
||||
transform_list = [TRANSFORM.build_with(pd) for pd in pipeline_description]
|
||||
return transforms.Compose(transform_list)
|
||||
|
||||
@@ -2,25 +2,27 @@ from itertools import chain
|
||||
|
||||
import ignite.distributed as idist
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from engine.base.i2i import EngineKernel, run_kernel
|
||||
from engine.util.build import build_model
|
||||
from loss.gan import GANLoss
|
||||
from model.GAN.base import GANImageBuffer
|
||||
from engine.util.container import GANImageBuffer, LossContainer
|
||||
from engine.util.loss import pixel_loss, gan_loss, feature_match_loss
|
||||
from loss.I2I.edge_loss import EdgeLoss
|
||||
from loss.I2I.minimal_geometry_distortion_constraint_loss import MGCLoss
|
||||
from model.weight_init import generation_init_weights
|
||||
|
||||
|
||||
class TAFGEngineKernel(EngineKernel):
|
||||
class CycleGANEngineKernel(EngineKernel):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
gan_loss_cfg = OmegaConf.to_container(config.loss.gan)
|
||||
gan_loss_cfg.pop("weight")
|
||||
self.gan_loss = GANLoss(**gan_loss_cfg).to(idist.device())
|
||||
self.cycle_loss = nn.L1Loss() if config.loss.cycle.level == 1 else nn.MSELoss()
|
||||
self.id_loss = nn.L1Loss() if config.loss.id.level == 1 else nn.MSELoss()
|
||||
self.gan_loss = gan_loss(config.loss.gan)
|
||||
self.cycle_loss = LossContainer(config.loss.cycle.weight, pixel_loss(config.loss.cycle.level))
|
||||
self.id_loss = LossContainer(config.loss.id.weight, pixel_loss(config.loss.id.level))
|
||||
self.mgc_loss = LossContainer(config.loss.mgc.weight, MGCLoss("opposite"))
|
||||
self.fm_loss = LossContainer(config.loss.fm.weight, feature_match_loss(1, "same"))
|
||||
self.edge_loss = LossContainer(config.loss.edge.weight, EdgeLoss(
|
||||
"HED", hed_pretrained_model_path=config.loss.edge.hed_pretrained_model_path).to(idist.device()))
|
||||
self.image_buffers = {k: GANImageBuffer(config.data.train.buffer_size or 50) for k in
|
||||
self.discriminators.keys()}
|
||||
|
||||
@@ -56,21 +58,23 @@ class TAFGEngineKernel(EngineKernel):
|
||||
images["b2a"] = self.generators["b2a"](batch["b"])
|
||||
images["a2b2a"] = self.generators["b2a"](images["a2b"])
|
||||
images["b2a2b"] = self.generators["a2b"](images["b2a"])
|
||||
if self.config.loss.id.weight > 0:
|
||||
if self.id_loss.weight > 0:
|
||||
images["a2a"] = self.generators["b2a"](batch["a"])
|
||||
images["b2b"] = self.generators["a2b"](batch["b"])
|
||||
return images
|
||||
|
||||
def criterion_generators(self, batch, generated) -> dict:
|
||||
loss = dict()
|
||||
for phase in ["a2b", "b2a"]:
|
||||
loss[f"cycle_{phase[0]}"] = self.config.loss.cycle.weight * self.cycle_loss(
|
||||
generated[f"{phase}2{phase[0]}"], batch[phase[0]])
|
||||
loss[f"gan_{phase}"] = self.config.loss.gan.weight * self.gan_loss(
|
||||
self.discriminators[phase[-1]](generated[phase]), True)
|
||||
if self.config.loss.id.weight > 0:
|
||||
loss[f"id_{phase[0]}"] = self.config.loss.id.weight * self.id_loss(
|
||||
generated[f"{phase[0]}2{phase[0]}"], batch[phase[0]])
|
||||
for ph in "ab":
|
||||
loss[f"cycle_{ph}"] = self.cycle_loss(generated["a2b2a" if ph == "a" else "b2a2b"], batch[ph])
|
||||
loss[f"id_{ph}"] = self.id_loss(generated[f"{ph}2{ph}"], batch[ph])
|
||||
loss[f"mgc_{ph}"] = self.mgc_loss(generated["a2b" if ph == "a" else "b2a"], batch[ph])
|
||||
prediction_fake = self.discriminators[ph](generated["a2b" if ph == "b" else "b2a"])
|
||||
loss[f"gan_{ph}"] = self.config.loss.gan.weight * self.gan_loss(prediction_fake, True)
|
||||
if self.fm_loss.weight > 0:
|
||||
prediction_real = self.discriminators[ph](batch[ph])
|
||||
loss[f"feature_match_{ph}"] = self.fm_loss(prediction_fake, prediction_real)
|
||||
loss[f"edge_{ph}"] = self.edge_loss(generated["a2b" if ph == "a" else "b2a"], batch[ph], gt_is_edge=False)
|
||||
return loss
|
||||
|
||||
def criterion_discriminators(self, batch, generated) -> dict:
|
||||
@@ -97,5 +101,5 @@ class TAFGEngineKernel(EngineKernel):
|
||||
|
||||
|
||||
def run(task, config, _):
|
||||
kernel = TAFGEngineKernel(config)
|
||||
kernel = CycleGANEngineKernel(config)
|
||||
run_kernel(task, config, kernel)
|
||||
86
engine/GauGAN.py
Normal file
86
engine/GauGAN.py
Normal file
@@ -0,0 +1,86 @@
|
||||
from itertools import chain
|
||||
|
||||
import torch
|
||||
|
||||
from engine.base.i2i import EngineKernel, run_kernel
|
||||
from engine.util.build import build_model
|
||||
from engine.util.container import GANImageBuffer, LossContainer
|
||||
from engine.util.loss import gan_loss, feature_match_loss, perceptual_loss
|
||||
from loss.I2I.minimal_geometry_distortion_constraint_loss import MGCLoss
|
||||
from model.weight_init import generation_init_weights
|
||||
|
||||
|
||||
class GauGANEngineKernel(EngineKernel):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
self.gan_loss = gan_loss(config.loss.gan)
|
||||
self.mgc_loss = LossContainer(config.loss.mgc.weight, MGCLoss("opposite"))
|
||||
self.fm_loss = LossContainer(config.loss.fm.weight, feature_match_loss(1, "same"))
|
||||
self.perceptual_loss = LossContainer(config.loss.perceptual.weight, perceptual_loss(config.loss.perceptual))
|
||||
|
||||
self.image_buffers = {k: GANImageBuffer(config.data.train.buffer_size or 50) for k in
|
||||
self.discriminators.keys()}
|
||||
|
||||
def build_models(self) -> (dict, dict):
|
||||
generators = dict(
|
||||
main=build_model(self.config.model.generator)
|
||||
)
|
||||
discriminators = dict(
|
||||
b=build_model(self.config.model.discriminator)
|
||||
)
|
||||
self.logger.debug(discriminators["b"])
|
||||
self.logger.debug(generators["main"])
|
||||
|
||||
for m in chain(generators.values(), discriminators.values()):
|
||||
generation_init_weights(m)
|
||||
|
||||
return generators, discriminators
|
||||
|
||||
def setup_after_g(self):
|
||||
for discriminator in self.discriminators.values():
|
||||
discriminator.requires_grad_(True)
|
||||
|
||||
def setup_before_g(self):
|
||||
for discriminator in self.discriminators.values():
|
||||
discriminator.requires_grad_(False)
|
||||
|
||||
def forward(self, batch, inference=False) -> dict:
|
||||
images = dict()
|
||||
with torch.set_grad_enabled(not inference):
|
||||
images["a2b"] = self.generators["main"](batch["a"])
|
||||
return images
|
||||
|
||||
def criterion_generators(self, batch, generated) -> dict:
|
||||
loss = dict()
|
||||
prediction_fake = self.discriminators["b"](generated["a2b"])
|
||||
loss["gan"] = self.config.loss.gan.weight * self.gan_loss(prediction_fake, True)
|
||||
loss["mgc"] = self.mgc_loss(generated["a2b"], batch["a"])
|
||||
loss["perceptual"] = self.perceptual_loss(generated["a2b"], batch["a"])
|
||||
if self.fm_loss.weight > 0:
|
||||
prediction_real = self.discriminators["b"](batch["b"])
|
||||
loss["feature_match"] = self.fm_loss(prediction_fake, prediction_real)
|
||||
return loss
|
||||
|
||||
def criterion_discriminators(self, batch, generated) -> dict:
|
||||
loss = dict()
|
||||
generated_image = self.image_buffers["b"].query(generated["a2b"].detach())
|
||||
loss["b"] = (self.gan_loss(self.discriminators["b"](generated_image), False, is_discriminator=True) +
|
||||
self.gan_loss(self.discriminators["b"](batch["b"]), True, is_discriminator=True)) / 2
|
||||
return loss
|
||||
|
||||
def intermediate_images(self, batch, generated) -> dict:
|
||||
"""
|
||||
returned dict must be like: {"a": [img1, img2, ...], "b": [img3, img4, ...]}
|
||||
:param batch:
|
||||
:param generated: dict of images
|
||||
:return: dict like: {"a": [img1, img2, ...], "b": [img3, img4, ...]}
|
||||
"""
|
||||
return dict(
|
||||
a=[batch["a"].detach(), generated["a2b"].detach()],
|
||||
)
|
||||
|
||||
|
||||
def run(task, config, _):
|
||||
kernel = GauGANEngineKernel(config)
|
||||
run_kernel(task, config, kernel)
|
||||
133
engine/TAFG.py
133
engine/TAFG.py
@@ -76,11 +76,14 @@ class TAFGEngineKernel(EngineKernel):
|
||||
contents["b"], styles["b"] = generator.encode(batch["b"]["edge"], batch["b"]["img"], "b", "b")
|
||||
for ph in "ab":
|
||||
images[f"{ph}2{ph}"] = generator.decode(contents[ph], styles[ph], ph)
|
||||
images["a2b"] = generator.decode(contents["a"], styles["b"], "b")
|
||||
contents["recon_a"], styles["recon_b"] = generator.encode(self.edge_loss.edge_extractor(images["a2b"]),
|
||||
images["a2b"], "b", "b")
|
||||
images["cycle_b"] = generator.decode(contents["b"], styles["recon_b"], "b")
|
||||
images["cycle_a"] = generator.decode(contents["recon_a"], styles["a"], "a")
|
||||
|
||||
if self.engine.state.epoch > self.config.misc.add_new_loss_epoch:
|
||||
styles[f"random_b"] = torch.randn_like(styles["b"]).to(idist.device())
|
||||
images["a2b"] = generator.decode(contents["a"], styles["random_b"], "b")
|
||||
contents["recon_a"], styles["recon_b"] = generator.encode(self.edge_loss.edge_extractor(images["a2b"]),
|
||||
images["a2b"], "b", "b")
|
||||
images["cycle_b"] = generator.decode(contents["b"], styles["recon_b"], "b")
|
||||
images["cycle_a"] = generator.decode(contents["recon_a"], styles["a"], "a")
|
||||
return dict(styles=styles, contents=contents, images=images)
|
||||
|
||||
def criterion_generators(self, batch, generated) -> dict:
|
||||
@@ -91,50 +94,76 @@ class TAFGEngineKernel(EngineKernel):
|
||||
loss[f"recon_image_{ph}"] = self.config.loss.recon.weight * self.recon_loss(
|
||||
generated["images"][f"{ph}2{ph}"], batch[ph]["img"])
|
||||
|
||||
pred_fake = self.discriminators[ph](generated["images"][f"a2{ph}"])
|
||||
pred_fake = self.discriminators[ph](generated["images"][f"{ph}2{ph}"])
|
||||
loss[f"gan_{ph}"] = 0
|
||||
for sub_pred_fake in pred_fake:
|
||||
# last output is actual prediction
|
||||
loss[f"gan_{ph}"] += self.gan_loss(sub_pred_fake[-1], True) * self.config.loss.gan.weight
|
||||
loss["recon_content_a"] = self.config.loss.content_recon.weight * self.content_recon_loss(
|
||||
generated["contents"]["a"], generated["contents"]["recon_a"]
|
||||
)
|
||||
loss["recon_style_b"] = self.config.loss.style_recon.weight * self.style_recon_loss(
|
||||
generated["styles"]["b"], generated["styles"]["recon_b"]
|
||||
)
|
||||
|
||||
if self.config.loss.perceptual.weight > 0:
|
||||
loss["perceptual_a"] = self.config.loss.perceptual.weight * self.perceptual_loss(
|
||||
batch["a"]["img"], generated["images"]["a2b"]
|
||||
if self.engine.state.epoch == self.config.misc.add_new_loss_epoch:
|
||||
self.generators["main"].style_converters.requires_grad_(False)
|
||||
self.generators["main"].style_encoders.requires_grad_(False)
|
||||
|
||||
if self.engine.state.epoch > self.config.misc.add_new_loss_epoch:
|
||||
|
||||
pred_fake = self.discriminators[ph](generated["images"]["a2b"])
|
||||
loss["gan_a2b"] = 0
|
||||
for sub_pred_fake in pred_fake:
|
||||
# last output is actual prediction
|
||||
loss["gan_a2b"] += self.gan_loss(sub_pred_fake[-1], True) * self.config.loss.gan.weight
|
||||
|
||||
loss["recon_content_a"] = self.config.loss.content_recon.weight * self.content_recon_loss(
|
||||
generated["contents"]["a"], generated["contents"]["recon_a"]
|
||||
)
|
||||
loss["recon_style_b"] = self.config.loss.style_recon.weight * self.style_recon_loss(
|
||||
generated["styles"]["random_b"], generated["styles"]["recon_b"]
|
||||
)
|
||||
|
||||
for ph in "ab":
|
||||
if self.config.loss.perceptual.weight > 0:
|
||||
loss["perceptual_a"] = self.config.loss.perceptual.weight * self.perceptual_loss(
|
||||
batch["a"]["img"], generated["images"]["a2b"]
|
||||
)
|
||||
|
||||
if self.config.loss.cycle.weight > 0:
|
||||
loss[f"cycle_{ph}"] = self.config.loss.cycle.weight * self.cycle_loss(
|
||||
batch[ph]["img"], generated["images"][f"cycle_{ph}"]
|
||||
)
|
||||
if self.config.loss.style.weight > 0:
|
||||
loss[f"style_{ph}"] = self.config.loss.style.weight * self.style_loss(
|
||||
batch[ph]["img"], generated["images"][f"a2{ph}"]
|
||||
loss[f"cycle_a"] = self.config.loss.cycle.weight * self.cycle_loss(
|
||||
batch["a"]["img"], generated["images"][f"cycle_a"]
|
||||
)
|
||||
|
||||
if self.config.loss.edge.weight > 0:
|
||||
loss["edge_a"] = self.config.loss.edge.weight * self.edge_loss(
|
||||
generated["images"]["a2b"], batch["a"]["edge"][:, 0:1, :, :]
|
||||
)
|
||||
# for ph in "ab":
|
||||
#
|
||||
# if self.config.loss.style.weight > 0:
|
||||
# loss[f"style_{ph}"] = self.config.loss.style.weight * self.style_loss(
|
||||
# batch[ph]["img"], generated["images"][f"a2{ph}"]
|
||||
# )
|
||||
|
||||
if self.config.loss.edge.weight > 0:
|
||||
loss["edge_a"] = self.config.loss.edge.weight * self.edge_loss(
|
||||
generated["images"]["a2b"], batch["a"]["edge"][:, 0:1, :, :]
|
||||
)
|
||||
|
||||
return loss
|
||||
|
||||
def criterion_discriminators(self, batch, generated) -> dict:
|
||||
loss = dict()
|
||||
# batch = self._process_batch(batch)
|
||||
for phase in self.discriminators.keys():
|
||||
pred_real = self.discriminators[phase](batch[phase]["img"])
|
||||
pred_fake = self.discriminators[phase](generated["images"][f"a2{phase}"].detach())
|
||||
loss[f"gan_{phase}"] = 0
|
||||
for i in range(len(pred_fake)):
|
||||
loss[f"gan_{phase}"] += (self.gan_loss(pred_fake[i][-1], False, is_discriminator=True)
|
||||
+ self.gan_loss(pred_real[i][-1], True, is_discriminator=True)) / 2
|
||||
|
||||
if self.engine.state.epoch > self.config.misc.add_new_loss_epoch:
|
||||
for phase in self.discriminators.keys():
|
||||
pred_real = self.discriminators[phase](batch[phase]["img"])
|
||||
pred_fake = self.discriminators[phase](generated["images"][f"{phase}2{phase}"].detach())
|
||||
pred_fake_2 = self.discriminators[phase](generated["images"]["a2b"].detach())
|
||||
loss[f"gan_{phase}"] = 0
|
||||
for i in range(len(pred_fake)):
|
||||
loss[f"gan_{phase}"] += (self.gan_loss(pred_fake[i][-1], False, is_discriminator=True) +
|
||||
self.gan_loss(pred_fake_2[i][-1], False, is_discriminator=True)
|
||||
+ self.gan_loss(pred_real[i][-1], True, is_discriminator=True)) / 3
|
||||
else:
|
||||
for phase in self.discriminators.keys():
|
||||
pred_real = self.discriminators[phase](batch[phase]["img"])
|
||||
pred_fake = self.discriminators[phase](generated["images"][f"{phase}2{phase}"].detach())
|
||||
loss[f"gan_{phase}"] = 0
|
||||
for i in range(len(pred_fake)):
|
||||
loss[f"gan_{phase}"] += (self.gan_loss(pred_fake[i][-1], False, is_discriminator=True)
|
||||
+ self.gan_loss(pred_real[i][-1], True, is_discriminator=True)) / 2
|
||||
return loss
|
||||
|
||||
def intermediate_images(self, batch, generated) -> dict:
|
||||
@@ -145,18 +174,30 @@ class TAFGEngineKernel(EngineKernel):
|
||||
:return: dict like: {"a": [img1, img2, ...], "b": [img3, img4, ...]}
|
||||
"""
|
||||
batch = self._process_batch(batch)
|
||||
return dict(
|
||||
a=[batch["a"]["edge"][:, 0:1, :, :].expand(-1, 3, -1, -1).detach(),
|
||||
batch["a"]["img"].detach(),
|
||||
generated["images"]["a2a"].detach(),
|
||||
generated["images"]["a2b"].detach(),
|
||||
generated["images"]["cycle_a"].detach(),
|
||||
],
|
||||
b=[batch["b"]["edge"].expand(-1, 3, -1, -1).detach(),
|
||||
batch["b"]["img"].detach(),
|
||||
generated["images"]["b2b"].detach(),
|
||||
generated["images"]["cycle_b"].detach()]
|
||||
)
|
||||
if self.engine.state.epoch > self.config.misc.add_new_loss_epoch:
|
||||
return dict(
|
||||
a=[batch["a"]["edge"][:, 0:1, :, :].expand(-1, 3, -1, -1).detach(),
|
||||
batch["a"]["img"].detach(),
|
||||
generated["images"]["a2a"].detach(),
|
||||
generated["images"]["a2b"].detach(),
|
||||
generated["images"]["cycle_a"].detach(),
|
||||
],
|
||||
b=[batch["b"]["edge"].expand(-1, 3, -1, -1).detach(),
|
||||
batch["b"]["img"].detach(),
|
||||
generated["images"]["b2b"].detach(),
|
||||
generated["images"]["cycle_b"].detach()]
|
||||
)
|
||||
else:
|
||||
return dict(
|
||||
a=[batch["a"]["edge"][:, 0:1, :, :].expand(-1, 3, -1, -1).detach(),
|
||||
batch["a"]["img"].detach(),
|
||||
generated["images"]["a2a"].detach(),
|
||||
],
|
||||
b=[batch["b"]["edge"].expand(-1, 3, -1, -1).detach(),
|
||||
batch["b"]["img"].detach(),
|
||||
generated["images"]["b2b"].detach(),
|
||||
]
|
||||
)
|
||||
|
||||
def change_engine(self, config, trainer):
|
||||
pass
|
||||
|
||||
@@ -51,31 +51,19 @@ class TSITEngineKernel(EngineKernel):
|
||||
def forward(self, batch, inference=False) -> dict:
|
||||
with torch.set_grad_enabled(not inference):
|
||||
fake = dict(
|
||||
b=self.generators["main"](content_img=batch["a"], style_img=batch["b"])
|
||||
b=self.generators["main"](content_img=batch["a"])
|
||||
)
|
||||
return fake
|
||||
|
||||
def criterion_generators(self, batch, generated) -> dict:
|
||||
loss = dict()
|
||||
loss_perceptual, _ = self.perceptual_loss(generated["b"], batch["a"])
|
||||
loss["perceptual"] = loss_perceptual * self.config.loss.perceptual.weight
|
||||
loss["perceptual"] = self.perceptual_loss(generated["b"], batch["a"]) * self.config.loss.perceptual.weight
|
||||
for phase in "b":
|
||||
pred_fake = self.discriminators[phase](generated[phase])
|
||||
loss[f"gan_{phase}"] = 0
|
||||
for sub_pred_fake in pred_fake:
|
||||
# last output is actual prediction
|
||||
loss[f"gan_{phase}"] += self.config.loss.gan.weight * self.gan_loss(sub_pred_fake[-1], True)
|
||||
|
||||
if self.config.loss.fm.weight > 0 and phase == "b":
|
||||
pred_real = self.discriminators[phase](batch[phase])
|
||||
loss_fm = 0
|
||||
num_scale_discriminator = len(pred_fake)
|
||||
for i in range(num_scale_discriminator):
|
||||
# last output is the final prediction, so we exclude it
|
||||
num_intermediate_outputs = len(pred_fake[i]) - 1
|
||||
for j in range(num_intermediate_outputs):
|
||||
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
|
||||
return loss
|
||||
|
||||
def criterion_discriminators(self, batch, generated) -> dict:
|
||||
|
||||
@@ -1,38 +1,38 @@
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import ignite.distributed as idist
|
||||
|
||||
from loss.gan import GANLoss
|
||||
from model.GAN.UGATIT import RhoClipper
|
||||
from model.GAN.base import GANImageBuffer
|
||||
from util.image import attention_colored_map
|
||||
from engine.base.i2i import EngineKernel, run_kernel, TestEngineKernel
|
||||
from engine.util.build import build_model
|
||||
from engine.util.container import LossContainer
|
||||
from engine.util.loss import bce_loss, mse_loss, pixel_loss, gan_loss
|
||||
from loss.I2I.minimal_geometry_distortion_constraint_loss import MyLoss
|
||||
from util.image import attention_colored_map
|
||||
|
||||
|
||||
def mse_loss(x, target_flag):
|
||||
return F.mse_loss(x, torch.ones_like(x) if target_flag else torch.zeros_like(x))
|
||||
class RhoClipper(object):
|
||||
def __init__(self, clip_min, clip_max):
|
||||
self.clip_min = clip_min
|
||||
self.clip_max = clip_max
|
||||
assert clip_min < clip_max
|
||||
|
||||
|
||||
def bce_loss(x, target_flag):
|
||||
return F.binary_cross_entropy_with_logits(x, torch.ones_like(x) if target_flag else torch.zeros_like(x))
|
||||
def __call__(self, module):
|
||||
if hasattr(module, 'rho'):
|
||||
w = module.rho.data
|
||||
w = w.clamp(self.clip_min, self.clip_max)
|
||||
module.rho.data = w
|
||||
|
||||
|
||||
class UGATITEngineKernel(EngineKernel):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
gan_loss_cfg = OmegaConf.to_container(config.loss.gan)
|
||||
gan_loss_cfg.pop("weight")
|
||||
self.gan_loss = GANLoss(**gan_loss_cfg).to(idist.device())
|
||||
self.cycle_loss = nn.L1Loss() if config.loss.cycle.level == 1 else nn.MSELoss()
|
||||
self.id_loss = nn.L1Loss() if config.loss.id.level == 1 else nn.MSELoss()
|
||||
self.gan_loss = gan_loss(config.loss.gan)
|
||||
self.cycle_loss = LossContainer(config.loss.cycle.weight, pixel_loss(config.loss.cycle.level))
|
||||
self.mgc_loss = LossContainer(config.loss.mgc.weight, MyLoss())
|
||||
self.id_loss = LossContainer(config.loss.id.weight, pixel_loss(config.loss.id.level))
|
||||
self.bce_loss = LossContainer(self.config.loss.cam.weight, bce_loss)
|
||||
self.mse_loss = LossContainer(self.config.loss.gan.weight, mse_loss)
|
||||
|
||||
self.rho_clipper = RhoClipper(0, 1)
|
||||
self.image_buffers = {k: GANImageBuffer(config.data.train.buffer_size or 50) for k in
|
||||
self.discriminators.keys()}
|
||||
self.train_generator_first = False
|
||||
|
||||
def build_models(self) -> (dict, dict):
|
||||
@@ -79,17 +79,17 @@ class UGATITEngineKernel(EngineKernel):
|
||||
loss = dict()
|
||||
for phase in "ab":
|
||||
cycle_image = generated["images"]["a2b2a" if phase == "a" else "b2a2b"]
|
||||
loss[f"cycle_{phase}"] = self.config.loss.cycle.weight * self.cycle_loss(cycle_image, batch[phase])
|
||||
loss[f"id_{phase}"] = self.config.loss.id.weight * self.id_loss(batch[phase],
|
||||
generated["images"][f"{phase}2{phase}"])
|
||||
loss[f"cycle_{phase}"] = self.cycle_loss(cycle_image, batch[phase])
|
||||
loss[f"id_{phase}"] = self.id_loss(batch[phase], generated["images"][f"{phase}2{phase}"])
|
||||
loss[f"mgc_{phase}"] = self.mgc_loss(batch[phase], generated["images"]["a2b" if phase == "a" else "b2a"])
|
||||
for dk in "lg":
|
||||
generated_image = generated["images"]["a2b" if phase == "b" else "b2a"]
|
||||
pred_fake, cam_pred = self.discriminators[dk + phase](generated_image)
|
||||
loss[f"gan_{phase}_{dk}"] = self.config.loss.gan.weight * self.gan_loss(pred_fake, True)
|
||||
loss[f"gan_cam_{phase}_{dk}"] = self.config.loss.gan.weight * mse_loss(cam_pred, True)
|
||||
loss[f"gan_cam_{phase}_{dk}"] = self.mse_loss(cam_pred, True)
|
||||
for t, f in [("a2b", "b2b"), ("b2a", "a2a")]:
|
||||
loss[f"cam_{t[-1]}"] = self.config.loss.cam.weight * (
|
||||
bce_loss(generated["cam_pred"][t], True) + bce_loss(generated["cam_pred"][f], False))
|
||||
loss[f"cam_{t[-1]}"] = self.bce_loss(generated["cam_pred"][t], True) + \
|
||||
self.bce_loss(generated["cam_pred"][f], False)
|
||||
return loss
|
||||
|
||||
def criterion_discriminators(self, batch, generated) -> dict:
|
||||
|
||||
@@ -64,7 +64,7 @@ class EngineKernel(object):
|
||||
self.engine = engine
|
||||
|
||||
def build_models(self) -> (dict, dict):
|
||||
raise NotImplemented
|
||||
raise NotImplementedError
|
||||
|
||||
def to_save(self):
|
||||
to_save = {}
|
||||
@@ -73,19 +73,19 @@ class EngineKernel(object):
|
||||
return to_save
|
||||
|
||||
def setup_after_g(self):
|
||||
raise NotImplemented
|
||||
raise NotImplementedError
|
||||
|
||||
def setup_before_g(self):
|
||||
raise NotImplemented
|
||||
raise NotImplementedError
|
||||
|
||||
def forward(self, batch, inference=False) -> dict:
|
||||
raise NotImplemented
|
||||
raise NotImplementedError
|
||||
|
||||
def criterion_generators(self, batch, generated) -> dict:
|
||||
raise NotImplemented
|
||||
raise NotImplementedError
|
||||
|
||||
def criterion_discriminators(self, batch, generated) -> dict:
|
||||
raise NotImplemented
|
||||
raise NotImplementedError
|
||||
|
||||
def intermediate_images(self, batch, generated) -> dict:
|
||||
"""
|
||||
@@ -94,12 +94,22 @@ class EngineKernel(object):
|
||||
:param generated: dict of images
|
||||
:return: dict like: {"a": [img1, img2, ...], "b": [img3, img4, ...]}
|
||||
"""
|
||||
raise NotImplemented
|
||||
raise NotImplementedError
|
||||
|
||||
def change_engine(self, config, engine: Engine):
|
||||
pass
|
||||
|
||||
|
||||
def _remove_no_grad_loss(loss_dict):
|
||||
need_to_pop = []
|
||||
for k in loss_dict:
|
||||
if not isinstance(loss_dict[k], torch.Tensor):
|
||||
need_to_pop.append(k)
|
||||
for k in need_to_pop:
|
||||
loss_dict.pop(k)
|
||||
return loss_dict
|
||||
|
||||
|
||||
def get_trainer(config, kernel: EngineKernel):
|
||||
logger = logging.getLogger(config.name)
|
||||
generators, discriminators = kernel.generators, kernel.discriminators
|
||||
@@ -147,10 +157,10 @@ def get_trainer(config, kernel: EngineKernel):
|
||||
|
||||
if engine.state.iteration % iteration_per_image == 0:
|
||||
return {
|
||||
"loss": dict(g=loss_g, d=loss_d),
|
||||
"loss": dict(g=_remove_no_grad_loss(loss_g), d=_remove_no_grad_loss(loss_d)),
|
||||
"img": kernel.intermediate_images(batch, generated)
|
||||
}
|
||||
return {"loss": dict(g=loss_g, d=loss_d)}
|
||||
return {"loss": dict(g=_remove_no_grad_loss(loss_g), d=_remove_no_grad_loss(loss_d))}
|
||||
|
||||
trainer = Engine(_step)
|
||||
trainer.logger = logger
|
||||
@@ -189,34 +199,33 @@ def get_trainer(config, kernel: EngineKernel):
|
||||
for i in range(len(image_list)):
|
||||
test_images[k].append([])
|
||||
|
||||
with torch.no_grad():
|
||||
g = torch.Generator()
|
||||
g.manual_seed(config.misc.random_seed + engine.state.epoch
|
||||
if config.handler.test.random else config.misc.random_seed)
|
||||
random_start = \
|
||||
torch.randperm(len(engine.state.test_dataset) - config.handler.test.images, generator=g).tolist()[0]
|
||||
for i in range(random_start, random_start + config.handler.test.images):
|
||||
batch = convert_tensor(engine.state.test_dataset[i], idist.device())
|
||||
for k in batch:
|
||||
if isinstance(batch[k], torch.Tensor):
|
||||
batch[k] = batch[k].unsqueeze(0)
|
||||
elif isinstance(batch[k], dict):
|
||||
for kk in batch[k]:
|
||||
if isinstance(batch[k][kk], torch.Tensor):
|
||||
batch[k][kk] = batch[k][kk].unsqueeze(0)
|
||||
g = torch.Generator()
|
||||
g.manual_seed(config.misc.random_seed + engine.state.epoch
|
||||
if config.handler.test.random else config.misc.random_seed)
|
||||
random_start = \
|
||||
torch.randperm(len(engine.state.test_dataset) - config.handler.test.images, generator=g).tolist()[0]
|
||||
for i in range(random_start, random_start + config.handler.test.images):
|
||||
batch = convert_tensor(engine.state.test_dataset[i], idist.device())
|
||||
for k in batch:
|
||||
if isinstance(batch[k], torch.Tensor):
|
||||
batch[k] = batch[k].unsqueeze(0)
|
||||
elif isinstance(batch[k], dict):
|
||||
for kk in batch[k]:
|
||||
if isinstance(batch[k][kk], torch.Tensor):
|
||||
batch[k][kk] = batch[k][kk].unsqueeze(0)
|
||||
|
||||
generated = kernel.forward(batch)
|
||||
images = kernel.intermediate_images(batch, generated)
|
||||
generated = kernel.forward(batch, inference=True)
|
||||
images = kernel.intermediate_images(batch, generated)
|
||||
|
||||
for k in test_images:
|
||||
for j in range(len(images[k])):
|
||||
test_images[k][j].append(images[k][j])
|
||||
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]], range=(-1, 1)),
|
||||
engine.state.iteration * pairs_per_iteration
|
||||
)
|
||||
for j in range(len(images[k])):
|
||||
test_images[k][j].append(images[k][j])
|
||||
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]], range=(-1, 1)),
|
||||
engine.state.iteration * pairs_per_iteration
|
||||
)
|
||||
return trainer
|
||||
|
||||
|
||||
|
||||
153
engine/talking_anime.py
Normal file
153
engine/talking_anime.py
Normal file
@@ -0,0 +1,153 @@
|
||||
from itertools import chain
|
||||
|
||||
import ignite.distributed as idist
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from engine.base.i2i import EngineKernel, run_kernel
|
||||
from engine.util.build import build_model
|
||||
from loss.I2I.context_loss import ContextLoss
|
||||
from loss.I2I.edge_loss import EdgeLoss
|
||||
from loss.I2I.perceptual_loss import PerceptualLoss
|
||||
from loss.gan import GANLoss
|
||||
from model.weight_init import generation_init_weights
|
||||
|
||||
|
||||
class TAEngineKernel(EngineKernel):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
|
||||
perceptual_loss_cfg = OmegaConf.to_container(config.loss.perceptual)
|
||||
perceptual_loss_cfg.pop("weight")
|
||||
self.perceptual_loss = PerceptualLoss(**perceptual_loss_cfg).to(idist.device())
|
||||
|
||||
style_loss_cfg = OmegaConf.to_container(config.loss.style)
|
||||
style_loss_cfg.pop("weight")
|
||||
self.style_loss = PerceptualLoss(**style_loss_cfg).to(idist.device())
|
||||
|
||||
gan_loss_cfg = OmegaConf.to_container(config.loss.gan)
|
||||
gan_loss_cfg.pop("weight")
|
||||
self.gan_loss = GANLoss(**gan_loss_cfg).to(idist.device())
|
||||
|
||||
context_loss_cfg = OmegaConf.to_container(config.loss.context)
|
||||
context_loss_cfg.pop("weight")
|
||||
self.context_loss = ContextLoss(**context_loss_cfg).to(idist.device())
|
||||
|
||||
self.recon_loss = nn.L1Loss() if config.loss.recon.level == 1 else nn.MSELoss()
|
||||
self.fm_loss = nn.L1Loss() if config.loss.fm.level == 1 else nn.MSELoss()
|
||||
|
||||
self.edge_loss = EdgeLoss("HED", hed_pretrained_model_path=config.loss.edge.hed_pretrained_model_path).to(
|
||||
idist.device())
|
||||
|
||||
def build_models(self) -> (dict, dict):
|
||||
generators = dict(
|
||||
anime=build_model(self.config.model.anime_generator),
|
||||
face=build_model(self.config.model.face_generator)
|
||||
)
|
||||
discriminators = dict(
|
||||
anime=build_model(self.config.model.discriminator),
|
||||
face=build_model(self.config.model.discriminator)
|
||||
)
|
||||
self.logger.debug(discriminators["face"])
|
||||
self.logger.debug(generators["face"])
|
||||
|
||||
for m in chain(generators.values(), discriminators.values()):
|
||||
generation_init_weights(m)
|
||||
|
||||
return generators, discriminators
|
||||
|
||||
def setup_after_g(self):
|
||||
for discriminator in self.discriminators.values():
|
||||
discriminator.requires_grad_(True)
|
||||
|
||||
def setup_before_g(self):
|
||||
for discriminator in self.discriminators.values():
|
||||
discriminator.requires_grad_(False)
|
||||
|
||||
def forward(self, batch, inference=False) -> dict:
|
||||
with torch.set_grad_enabled(not inference):
|
||||
target_pose_anime = self.generators["anime"](
|
||||
torch.cat([batch["face_1"], torch.flip(batch["anime_img"], dims=[3])], dim=1))
|
||||
target_pose_face = self.generators["face"](target_pose_anime.mean(dim=1, keepdim=True), batch["face_0"])
|
||||
|
||||
return dict(fake_anime=target_pose_anime, fake_face=target_pose_face)
|
||||
|
||||
def cal_gan_and_fm_loss(self, discriminator, generated_img, match_img=None):
|
||||
pred_fake = discriminator(generated_img)
|
||||
loss_gan = 0
|
||||
for sub_pred_fake in pred_fake:
|
||||
# last output is actual prediction
|
||||
loss_gan += self.config.loss.gan.weight * self.gan_loss(sub_pred_fake[-1], True)
|
||||
|
||||
if match_img is None:
|
||||
# do not cal feature match loss
|
||||
return loss_gan, 0
|
||||
|
||||
pred_real = discriminator(match_img)
|
||||
loss_fm = 0
|
||||
num_scale_discriminator = len(pred_fake)
|
||||
for i in range(num_scale_discriminator):
|
||||
# last output is the final prediction, so we exclude it
|
||||
num_intermediate_outputs = len(pred_fake[i]) - 1
|
||||
for j in range(num_intermediate_outputs):
|
||||
loss_fm += self.fm_loss(pred_fake[i][j], pred_real[i][j].detach()) / num_scale_discriminator
|
||||
loss_fm = self.config.loss.fm.weight * loss_fm
|
||||
return loss_gan, loss_fm
|
||||
|
||||
def criterion_generators(self, batch, generated) -> dict:
|
||||
loss = dict()
|
||||
loss["face_style"] = self.config.loss.style.weight * self.style_loss(
|
||||
generated["fake_face"], batch["face_1"]
|
||||
)
|
||||
loss["face_recon"] = self.config.loss.recon.weight * self.recon_loss(
|
||||
generated["fake_face"], batch["face_1"]
|
||||
)
|
||||
loss["face_gan"], loss["face_fm"] = self.cal_gan_and_fm_loss(
|
||||
self.discriminators["face"], generated["fake_face"], batch["face_1"])
|
||||
loss["anime_gan"], loss["anime_fm"] = self.cal_gan_and_fm_loss(
|
||||
self.discriminators["anime"], generated["fake_anime"], batch["anime_img"])
|
||||
|
||||
loss["anime_edge"] = self.config.loss.edge.weight * self.edge_loss(
|
||||
generated["fake_anime"], batch["face_1"], gt_is_edge=False,
|
||||
)
|
||||
if self.config.loss.perceptual.weight > 0:
|
||||
loss["anime_perceptual"] = self.config.loss.perceptual.weight * self.perceptual_loss(
|
||||
generated["fake_anime"], batch["anime_img"]
|
||||
)
|
||||
if self.config.loss.context.weight > 0:
|
||||
loss["anime_context"] = self.config.loss.context.weight * self.context_loss(
|
||||
generated["fake_anime"], batch["anime_img"],
|
||||
)
|
||||
|
||||
return loss
|
||||
|
||||
def criterion_discriminators(self, batch, generated) -> dict:
|
||||
loss = dict()
|
||||
real = {"anime": "anime_img", "face": "face_1"}
|
||||
for phase in self.discriminators.keys():
|
||||
pred_real = self.discriminators[phase](batch[real[phase]])
|
||||
pred_fake = self.discriminators[phase](generated[f"fake_{phase}"].detach())
|
||||
loss[f"gan_{phase}"] = 0
|
||||
for i in range(len(pred_fake)):
|
||||
loss[f"gan_{phase}"] += (self.gan_loss(pred_fake[i][-1], False, is_discriminator=True)
|
||||
+ self.gan_loss(pred_real[i][-1], True, is_discriminator=True)) / 2
|
||||
return loss
|
||||
|
||||
def intermediate_images(self, batch, generated) -> dict:
|
||||
"""
|
||||
returned dict must be like: {"a": [img1, img2, ...], "b": [img3, img4, ...]}
|
||||
:param batch:
|
||||
:param generated: dict of images
|
||||
:return: dict like: {"a": [img1, img2, ...], "b": [img3, img4, ...]}
|
||||
"""
|
||||
images = [batch["face_0"], batch["face_1"], batch["anime_img"], generated["fake_anime"].detach(),
|
||||
generated["fake_face"].detach()]
|
||||
return dict(
|
||||
b=[img for img in images]
|
||||
)
|
||||
|
||||
|
||||
def run(task, config, _):
|
||||
kernel = TAEngineKernel(config)
|
||||
run_kernel(task, config, kernel)
|
||||
@@ -1,18 +1,28 @@
|
||||
import torch
|
||||
import ignite.distributed as idist
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from model import MODEL
|
||||
import torch.optim as optim
|
||||
|
||||
|
||||
def add_spectral_norm(module):
|
||||
if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)) and not hasattr(module, 'weight_u'):
|
||||
return nn.utils.spectral_norm(module)
|
||||
else:
|
||||
return module
|
||||
|
||||
|
||||
def build_model(cfg):
|
||||
cfg = OmegaConf.to_container(cfg)
|
||||
bn_to_sync_bn = cfg.pop("_bn_to_sync_bn", False)
|
||||
add_spectral_norm_flag = cfg.pop("_add_spectral_norm", False)
|
||||
model = MODEL.build_with(cfg)
|
||||
if bn_to_sync_bn:
|
||||
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
|
||||
if add_spectral_norm_flag:
|
||||
model.apply(add_spectral_norm)
|
||||
return idist.auto_model(model)
|
||||
|
||||
|
||||
|
||||
66
engine/util/container.py
Normal file
66
engine/util/container.py
Normal file
@@ -0,0 +1,66 @@
|
||||
import torch
|
||||
|
||||
|
||||
class LossContainer:
|
||||
def __init__(self, weight, loss):
|
||||
self.weight = weight
|
||||
self.loss = loss
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
if self.weight > 0:
|
||||
return self.weight * self.loss(*args, **kwargs)
|
||||
return 0.0
|
||||
|
||||
|
||||
class GANImageBuffer:
|
||||
"""This class implements an image buffer that stores previously
|
||||
generated images.
|
||||
This buffer allows us to update the discriminator using a history of
|
||||
generated images rather than the ones produced by the latest generator
|
||||
to reduce model oscillation.
|
||||
Args:
|
||||
buffer_size (int): The size of image buffer. If buffer_size = 0,
|
||||
no buffer will be created.
|
||||
buffer_ratio (float): The chance / possibility to use the images
|
||||
previously stored in the buffer.
|
||||
"""
|
||||
|
||||
def __init__(self, buffer_size, buffer_ratio=0.5):
|
||||
self.buffer_size = buffer_size
|
||||
# create an empty buffer
|
||||
if self.buffer_size > 0:
|
||||
self.img_num = 0
|
||||
self.image_buffer = []
|
||||
self.buffer_ratio = buffer_ratio
|
||||
|
||||
def query(self, images):
|
||||
"""Query current image batch using a history of generated images.
|
||||
Args:
|
||||
images (Tensor): Current image batch without history information.
|
||||
"""
|
||||
if self.buffer_size == 0: # if the buffer size is 0, do nothing
|
||||
return images
|
||||
return_images = []
|
||||
for image in images:
|
||||
image = torch.unsqueeze(image.data, 0)
|
||||
# if the buffer is not full, keep inserting current images
|
||||
if self.img_num < self.buffer_size:
|
||||
self.img_num = self.img_num + 1
|
||||
self.image_buffer.append(image)
|
||||
return_images.append(image)
|
||||
else:
|
||||
use_buffer = torch.rand(1) < self.buffer_ratio
|
||||
# by self.buffer_ratio, the buffer will return a previously
|
||||
# stored image, and insert the current image into the buffer
|
||||
if use_buffer:
|
||||
random_id = torch.randint(0, self.buffer_size, (1,)).item()
|
||||
image_tmp = self.image_buffer[random_id].clone()
|
||||
self.image_buffer[random_id] = image
|
||||
return_images.append(image_tmp)
|
||||
# by (1 - self.buffer_ratio), the buffer will return the
|
||||
# current image
|
||||
else:
|
||||
return_images.append(image)
|
||||
# collect all the images and return
|
||||
return_images = torch.cat(return_images, 0)
|
||||
return return_images
|
||||
48
engine/util/loss.py
Normal file
48
engine/util/loss.py
Normal file
@@ -0,0 +1,48 @@
|
||||
import ignite.distributed as idist
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from loss.I2I.perceptual_loss import PerceptualLoss
|
||||
from loss.gan import GANLoss
|
||||
|
||||
|
||||
def gan_loss(config):
|
||||
gan_loss_cfg = OmegaConf.to_container(config)
|
||||
gan_loss_cfg.pop("weight")
|
||||
return GANLoss(**gan_loss_cfg).to(idist.device())
|
||||
|
||||
|
||||
def perceptual_loss(config):
|
||||
perceptual_loss_cfg = OmegaConf.to_container(config)
|
||||
perceptual_loss_cfg.pop("weight")
|
||||
return PerceptualLoss(**perceptual_loss_cfg).to(idist.device())
|
||||
|
||||
|
||||
def pixel_loss(level):
|
||||
return nn.L1Loss() if level == 1 else nn.MSELoss()
|
||||
|
||||
|
||||
def mse_loss(x, target_flag):
|
||||
return F.mse_loss(x, torch.ones_like(x) if target_flag else torch.zeros_like(x))
|
||||
|
||||
|
||||
def bce_loss(x, target_flag):
|
||||
return F.binary_cross_entropy_with_logits(x, torch.ones_like(x) if target_flag else torch.zeros_like(x))
|
||||
|
||||
|
||||
def feature_match_loss(level, weight_policy):
|
||||
compare_loss = pixel_loss(level)
|
||||
assert weight_policy in ["same", "exponential_decline"]
|
||||
|
||||
def fm_loss(generated_features, target_features):
|
||||
num_scale = len(generated_features)
|
||||
loss = torch.zeros(1, device=idist.device())
|
||||
for s_i in range(num_scale):
|
||||
for i in range(len(generated_features[s_i]) - 1):
|
||||
weight = 1 if weight_policy == "same" else 2 ** i
|
||||
loss += weight * compare_loss(generated_features[s_i][i], target_features[s_i][i].detach()) / num_scale
|
||||
return loss
|
||||
|
||||
return fm_loss
|
||||
44
loss/I2I/context_loss.py
Normal file
44
loss/I2I/context_loss.py
Normal file
@@ -0,0 +1,44 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from .perceptual_loss import PerceptualVGG
|
||||
|
||||
|
||||
class ContextLoss(nn.Module):
|
||||
def __init__(self, layer_weights, h=0.1, vgg_type='vgg19', norm_image_with_imagenet_param=True, norm_img=True,
|
||||
eps=1e-5):
|
||||
super(ContextLoss, self).__init__()
|
||||
self.eps = eps
|
||||
self.h = h
|
||||
self.layer_weights = layer_weights
|
||||
self.norm_img = norm_img
|
||||
self.vgg = PerceptualVGG(layer_name_list=list(layer_weights.keys()), vgg_type=vgg_type,
|
||||
norm_image_with_imagenet_param=norm_image_with_imagenet_param)
|
||||
|
||||
def single_forward(self, source_feature, target_feature):
|
||||
mean_target_feature = target_feature.mean(dim=[2, 3], keepdim=True)
|
||||
source_feature = (source_feature - mean_target_feature).view(*source_feature.size()[:2], -1) # NxCxHW
|
||||
target_feature = (target_feature - mean_target_feature).view(*source_feature.size()[:2], -1) # NxCxHW
|
||||
source_feature = F.normalize(source_feature, p=2, dim=1)
|
||||
target_feature = F.normalize(target_feature, p=2, dim=1)
|
||||
cosine_distance = (1 - torch.bmm(source_feature.transpose(1, 2), target_feature)) / 2 # NxHWxHW
|
||||
rel_distance = cosine_distance / (cosine_distance.min(2, keepdim=True)[0] + self.eps)
|
||||
w = torch.exp((1 - rel_distance) / self.h)
|
||||
cx = w.div(w.sum(dim=2, keepdim=True))
|
||||
cx = cx.max(dim=1, keepdim=True)[0].mean(dim=2)
|
||||
return -torch.log(cx).mean()
|
||||
|
||||
def forward(self, x, gt):
|
||||
if self.norm_img:
|
||||
x = (x + 1.) * 0.5
|
||||
gt = (gt + 1.) * 0.5
|
||||
|
||||
# extract vgg features
|
||||
x_features = self.vgg(x)
|
||||
gt_features = self.vgg(gt.detach())
|
||||
|
||||
loss = 0
|
||||
for k in x_features.keys():
|
||||
loss += self.single_forward(x_features[k], gt_features[k]) * self.layer_weights[k]
|
||||
return loss
|
||||
229
loss/I2I/minimal_geometry_distortion_constraint_loss.py
Normal file
229
loss/I2I/minimal_geometry_distortion_constraint_loss.py
Normal file
@@ -0,0 +1,229 @@
|
||||
import ignite.distributed as idist
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def gaussian_radial_basis_function(x, mu, sigma):
|
||||
# (kernel_size) -> (batch_size, kernel_size, c*h*w)
|
||||
mu = mu.view(1, mu.size(0), 1).expand(x.size(0), -1, x.size(1) * x.size(2) * x.size(3))
|
||||
# (batch_size, c, h, w) -> (batch_size, kernel_size, c*h*w)
|
||||
x = x.view(x.size(0), 1, -1).expand(-1, mu.size(1), -1)
|
||||
return torch.exp((x - mu).pow(2) / (2 * sigma ** 2))
|
||||
|
||||
|
||||
class ImporveMyLoss(torch.nn.Module):
|
||||
def __init__(self, device=idist.device()):
|
||||
super().__init__()
|
||||
mu = torch.Tensor([-1.0, -0.75, -0.5, -0.25, 0, 0.25, 0.5, 0.75, 1.0]).to(device)
|
||||
self.x_mu_list = mu.repeat(9).view(-1, 81)
|
||||
self.y_mu_list = mu.unsqueeze(0).t().repeat(1, 9).view(-1, 81)
|
||||
self.R = torch.eye(81).to(device)
|
||||
|
||||
def batch_ERSMI(self, I1, I2):
|
||||
batch_size = I1.shape[0]
|
||||
img_size = I1.shape[1] * I1.shape[2] * I1.shape[3]
|
||||
if I2.shape[1] == 1 and I1.shape[1] != 1:
|
||||
I2 = I2.repeat(1, 3, 1, 1)
|
||||
|
||||
def kernel_F(y, mu_list, sigma):
|
||||
tmp_mu = mu_list.view(-1, 1).repeat(1, img_size).repeat(batch_size, 1, 1) # [81, 784]
|
||||
tmp_y = y.view(batch_size, 1, -1).repeat(1, 81, 1)
|
||||
tmp_y = tmp_mu - tmp_y
|
||||
mat_L = torch.exp(tmp_y.pow(2) / (2 * sigma ** 2))
|
||||
return mat_L
|
||||
|
||||
mat_K = kernel_F(I1, self.x_mu_list, 1)
|
||||
mat_L = kernel_F(I2, self.y_mu_list, 1)
|
||||
mat_k_l = mat_K * mat_L
|
||||
|
||||
H1 = (mat_K @ mat_K.transpose(1, 2)) * (mat_L @ mat_L.transpose(1, 2)) / (img_size ** 2)
|
||||
h_hat = mat_k_l @ mat_k_l.transpose(1, 2) / img_size
|
||||
small_h_hat = mat_K.sum(2).view(batch_size, -1, 1) * mat_L.sum(2).view(batch_size, -1, 1) / (img_size ** 2)
|
||||
h_hat = 0.5 * H1 + 0.5 * h_hat
|
||||
alpha = (h_hat + 0.05 * self.R).inverse() @ small_h_hat
|
||||
|
||||
ersmi = 2 * alpha.transpose(1, 2) @ small_h_hat - alpha.transpose(1, 2) @ h_hat @ alpha - 1
|
||||
|
||||
ersmi = -ersmi.squeeze().mean()
|
||||
return ersmi
|
||||
|
||||
def forward(self, fakeI, realI):
|
||||
return self.batch_ERSMI(fakeI, realI)
|
||||
|
||||
|
||||
class MyLoss(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super(MyLoss, self).__init__()
|
||||
|
||||
def forward(self, fakeI, realI):
|
||||
fakeI = fakeI.cuda()
|
||||
realI = realI.cuda()
|
||||
|
||||
def batch_ERSMI(I1, I2):
|
||||
batch_size = I1.shape[0]
|
||||
img_size = I1.shape[1] * I1.shape[2] * I1.shape[3]
|
||||
if I2.shape[1] == 1 and I1.shape[1] != 1:
|
||||
I2 = I2.repeat(1, 3, 1, 1)
|
||||
|
||||
def kernel_F(y, mu_list, sigma):
|
||||
tmp_mu = mu_list.view(-1, 1).repeat(1, img_size).repeat(batch_size, 1, 1).cuda() # [81, 784]
|
||||
tmp_y = y.view(batch_size, 1, -1).repeat(1, 81, 1)
|
||||
tmp_y = tmp_mu - tmp_y
|
||||
mat_L = torch.exp(tmp_y.pow(2) / (2 * sigma ** 2))
|
||||
return mat_L
|
||||
|
||||
mu = torch.Tensor([-1.0, -0.75, -0.5, -0.25, 0, 0.25, 0.5, 0.75, 1.0]).cuda()
|
||||
|
||||
x_mu_list = mu.repeat(9).view(-1, 81)
|
||||
y_mu_list = mu.unsqueeze(0).t().repeat(1, 9).view(-1, 81)
|
||||
|
||||
mat_K = kernel_F(I1, x_mu_list, 1)
|
||||
mat_L = kernel_F(I2, y_mu_list, 1)
|
||||
|
||||
H1 = ((mat_K.matmul(mat_K.transpose(1, 2))).mul(mat_L.matmul(mat_L.transpose(1, 2))) / (
|
||||
img_size ** 2)).cuda()
|
||||
H2 = ((mat_K.mul(mat_L)).matmul((mat_K.mul(mat_L)).transpose(1, 2)) / img_size).cuda()
|
||||
h2 = ((mat_K.sum(2).view(batch_size, -1, 1)).mul(mat_L.sum(2).view(batch_size, -1, 1)) / (
|
||||
img_size ** 2)).cuda()
|
||||
H2 = 0.5 * H1 + 0.5 * H2
|
||||
tmp = H2 + 0.05 * torch.eye(len(H2[0])).cuda()
|
||||
alpha = (tmp.inverse())
|
||||
|
||||
alpha = alpha.matmul(h2)
|
||||
ersmi = (2 * (alpha.transpose(1, 2)).matmul(h2) - ((alpha.transpose(1, 2)).matmul(H2)).matmul(
|
||||
alpha) - 1).squeeze()
|
||||
|
||||
ersmi = -ersmi.mean()
|
||||
return ersmi
|
||||
|
||||
batch_loss = batch_ERSMI(fakeI, realI)
|
||||
return batch_loss
|
||||
|
||||
|
||||
class MGCLoss(nn.Module):
|
||||
"""
|
||||
Minimal Geometry-Distortion Constraint Loss from https://openreview.net/forum?id=R5M7Mxl1xZ
|
||||
"""
|
||||
|
||||
def __init__(self, mi_to_loss_way="opposite", beta=0.5, lambda_=0.05, device=idist.device()):
|
||||
super().__init__()
|
||||
self.beta = beta
|
||||
self.lambda_ = lambda_
|
||||
assert mi_to_loss_way in ["opposite", "reciprocal"]
|
||||
self.mi_to_loss_way = mi_to_loss_way
|
||||
mu_y, mu_x = torch.meshgrid([torch.arange(-1, 1.25, 0.25), torch.arange(-1, 1.25, 0.25)])
|
||||
self.mu_x = mu_x.flatten().to(device)
|
||||
self.mu_y = mu_y.flatten().to(device)
|
||||
self.R = torch.eye(81).unsqueeze(0).to(device)
|
||||
|
||||
@staticmethod
|
||||
def batch_rSMI(img1, img2, mu_x, mu_y, beta, lambda_, R):
|
||||
assert img1.size() == img2.size()
|
||||
|
||||
num_pixel = img1.size(1) * img1.size(2) * img2.size(3)
|
||||
|
||||
mat_k = gaussian_radial_basis_function(img1, mu_x, sigma=1)
|
||||
mat_l = gaussian_radial_basis_function(img2, mu_y, sigma=1)
|
||||
|
||||
mat_k_mul_mat_l = mat_k * mat_l
|
||||
h_hat = (1 - beta) * (mat_k_mul_mat_l @ mat_k_mul_mat_l.transpose(1, 2)) / num_pixel
|
||||
h_hat += beta * ((mat_k @ mat_k.transpose(1, 2)) * (mat_l @ mat_l.transpose(1, 2))) / (num_pixel ** 2)
|
||||
small_h_hat = mat_k.sum(2, keepdim=True) * mat_l.sum(2, keepdim=True) / (num_pixel ** 2)
|
||||
|
||||
alpha = (h_hat + lambda_ * R).inverse() @ small_h_hat
|
||||
rSMI = 2 * alpha.transpose(1, 2) @ small_h_hat - alpha.transpose(1, 2) @ h_hat @ alpha - 1
|
||||
return rSMI.squeeze()
|
||||
|
||||
def forward(self, fake, real):
|
||||
rSMI = self.batch_rSMI(fake, real, self.mu_x, self.mu_y, self.beta, self.lambda_, self.R)
|
||||
if self.mi_to_loss_way == "reciprocal":
|
||||
return 1/rSMI.mean()
|
||||
return -rSMI.mean()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
mg = MGCLoss(device=torch.device("cpu"))
|
||||
my = MyLoss().to("cuda")
|
||||
imy = ImporveMyLoss()
|
||||
|
||||
from data.transform import transform_pipeline
|
||||
|
||||
pipeline = transform_pipeline(
|
||||
['Load', 'ToTensor', {'Normalize': {'mean': [0.5, 0.5, 0.5], 'std': [0.5, 0.5, 0.5]}}])
|
||||
|
||||
img_a1 = pipeline("/data/i2i/VoxCeleb2Anime/trainA/id00022-twCPGo2rtCo-00294_1.jpg")
|
||||
img_a2 = pipeline("/data/i2i/VoxCeleb2Anime/trainA/id00022-twCPGo2rtCo-00294_2.jpg")
|
||||
img_a3 = pipeline("/data/i2i/VoxCeleb2Anime/trainA/id00022-twCPGo2rtCo-00294_3.jpg")
|
||||
img_b1 = pipeline("/data/i2i/VoxCeleb2Anime/trainA/id01222-2gHw81dNQiA-00005_1.jpg")
|
||||
img_b2 = pipeline("/data/i2i/VoxCeleb2Anime/trainA/id01222-2gHw81dNQiA-00005_2.jpg")
|
||||
img_b3 = pipeline("/data/i2i/VoxCeleb2Anime/trainA/id01222-2gHw81dNQiA-00005_3.jpg")
|
||||
|
||||
img_a1.requires_grad_(True)
|
||||
img_a2.requires_grad_(True)
|
||||
img_a3.requires_grad_(True)
|
||||
|
||||
# print("MyLoss")
|
||||
# l1 = my(img_a1.unsqueeze(0), img_b1.unsqueeze(0))
|
||||
# l2 = my(img_a2.unsqueeze(0), img_b2.unsqueeze(0))
|
||||
# l3 = my(img_a3.unsqueeze(0), img_b3.unsqueeze(0))
|
||||
# l = (l1+l2+l3)/3
|
||||
# l.backward()
|
||||
# print(img_a1.grad[0][0][0:10])
|
||||
# print(img_a2.grad[0][0][0:10])
|
||||
# print(img_a3.grad[0][0][0:10])
|
||||
#
|
||||
# img_a1.grad = None
|
||||
# img_a2.grad = None
|
||||
# img_a3.grad = None
|
||||
#
|
||||
# print("---")
|
||||
# l = my(torch.stack([img_a1, img_a2, img_a3]), torch.stack([img_b1, img_b2, img_b3]))
|
||||
# l.backward()
|
||||
# print(img_a1.grad[0][0][0:10])
|
||||
# print(img_a2.grad[0][0][0:10])
|
||||
# print(img_a3.grad[0][0][0:10])
|
||||
# img_a1.grad = None
|
||||
# img_a2.grad = None
|
||||
# img_a3.grad = None
|
||||
|
||||
print("MGCLoss")
|
||||
l1 = mg(img_a1.unsqueeze(0), img_b1.unsqueeze(0))
|
||||
l2 = mg(img_a2.unsqueeze(0), img_b2.unsqueeze(0))
|
||||
l3 = mg(img_a3.unsqueeze(0), img_b3.unsqueeze(0))
|
||||
l = (l1 + l2 + l3) / 3
|
||||
l.backward()
|
||||
print(img_a1.grad[0][0][0:10])
|
||||
print(img_a2.grad[0][0][0:10])
|
||||
print(img_a3.grad[0][0][0:10])
|
||||
|
||||
img_a1.grad = None
|
||||
img_a2.grad = None
|
||||
img_a3.grad = None
|
||||
|
||||
print("---")
|
||||
l = mg(torch.stack([img_a1, img_a2, img_a3]), torch.stack([img_b1, img_b2, img_b3]))
|
||||
l.backward()
|
||||
print(img_a1.grad[0][0][0:10])
|
||||
print(img_a2.grad[0][0][0:10])
|
||||
print(img_a3.grad[0][0][0:10])
|
||||
|
||||
# print("\nMGCLoss")
|
||||
# mg(img_a1.unsqueeze(0), img_b1.unsqueeze(0))
|
||||
# mg(img_a2.unsqueeze(0), img_b2.unsqueeze(0))
|
||||
# mg(img_a3.unsqueeze(0), img_b3.unsqueeze(0))
|
||||
#
|
||||
# print("---")
|
||||
# mg(torch.stack([img_a1, img_a2, img_a3]), torch.stack([img_b1, img_b2, img_b3]))
|
||||
#
|
||||
# import pprofile
|
||||
#
|
||||
# profiler = pprofile.Profile()
|
||||
# with profiler:
|
||||
# iter_times = 1000
|
||||
# for _ in range(iter_times):
|
||||
# mg(torch.stack([img_a1, img_a2, img_a3]), torch.stack([img_b1, img_b2, img_b3]))
|
||||
# for _ in range(iter_times):
|
||||
# my(torch.stack([img_a1, img_a2, img_a3]), torch.stack([img_b1, img_b2, img_b3]))
|
||||
# for _ in range(iter_times):
|
||||
# imy(torch.stack([img_a1, img_a2, img_a3]), torch.stack([img_b1, img_b2, img_b3]))
|
||||
# profiler.print_stats()
|
||||
@@ -4,6 +4,49 @@ import torch.nn.functional as F
|
||||
import torchvision.models.vgg as vgg
|
||||
|
||||
|
||||
# Sequential(
|
||||
# (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (1): ReLU(inplace=True)
|
||||
# (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (3): ReLU(inplace=True)
|
||||
# (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||||
|
||||
# (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (6): ReLU(inplace=True)
|
||||
# (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (8): ReLU(inplace=True)
|
||||
# (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||||
|
||||
# (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (11): ReLU(inplace=True)
|
||||
# (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (13): ReLU(inplace=True)
|
||||
# (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (15): ReLU(inplace=True)
|
||||
# (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (17): ReLU(inplace=True)
|
||||
# (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||||
|
||||
# (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (20): ReLU(inplace=True)
|
||||
# (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (22): ReLU(inplace=True)
|
||||
# (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (24): ReLU(inplace=True)
|
||||
# (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (26): ReLU(inplace=True)
|
||||
# (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||||
|
||||
# (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (29): ReLU(inplace=True)
|
||||
# (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (31): ReLU(inplace=True)
|
||||
# (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (33): ReLU(inplace=True)
|
||||
# (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
|
||||
# (35): ReLU(inplace=True)
|
||||
# (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
|
||||
# )
|
||||
class PerceptualVGG(nn.Module):
|
||||
"""VGG network used in calculating perceptual loss.
|
||||
In this implementation, we allow users to choose whether use normalization
|
||||
@@ -15,15 +58,15 @@ class PerceptualVGG(nn.Module):
|
||||
list contains the name each layer in `vgg.feature`. An example
|
||||
of this list is ['4', '10'].
|
||||
vgg_type (str): Set the type of vgg network. Default: 'vgg19'.
|
||||
use_input_norm (bool): If True, normalize the input image.
|
||||
norm_image_with_imagenet_param (bool): If True, normalize the input image.
|
||||
Importantly, the input feature must in the range [0, 1].
|
||||
Default: True.
|
||||
"""
|
||||
|
||||
def __init__(self, layer_name_list, vgg_type='vgg19', use_input_norm=True):
|
||||
def __init__(self, layer_name_list, vgg_type='vgg19', norm_image_with_imagenet_param=True):
|
||||
super(PerceptualVGG, self).__init__()
|
||||
self.layer_name_list = layer_name_list
|
||||
self.use_input_norm = use_input_norm
|
||||
self.use_input_norm = norm_image_with_imagenet_param
|
||||
|
||||
# get vgg model and load pretrained vgg weight
|
||||
# remove _vgg from attributes to avoid `find_unused_parameters` bug
|
||||
@@ -75,7 +118,7 @@ class PerceptualLoss(nn.Module):
|
||||
in calculating losses.
|
||||
vgg_type (str): The type of vgg network used as feature extractor.
|
||||
Default: 'vgg19'.
|
||||
use_input_norm (bool): If True, normalize the input image in vgg.
|
||||
norm_image_with_imagenet_param (bool): If True, normalize the input image in vgg.
|
||||
Default: True.
|
||||
perceptual_loss (bool): If `perceptual_loss == True`, the perceptual
|
||||
loss will be calculated.
|
||||
@@ -88,7 +131,7 @@ class PerceptualLoss(nn.Module):
|
||||
Importantly, the input image must be in range [-1, 1].
|
||||
"""
|
||||
|
||||
def __init__(self, layer_weights, vgg_type='vgg19', use_input_norm=True, perceptual_loss=True,
|
||||
def __init__(self, layer_weights, vgg_type='vgg19', norm_image_with_imagenet_param=True, perceptual_loss=True,
|
||||
style_loss=False, norm_img=True, criterion='L1'):
|
||||
super(PerceptualLoss, self).__init__()
|
||||
self.norm_img = norm_img
|
||||
@@ -97,7 +140,7 @@ class PerceptualLoss(nn.Module):
|
||||
self.style_loss = style_loss
|
||||
self.layer_weights = layer_weights
|
||||
self.vgg = PerceptualVGG(layer_name_list=list(layer_weights.keys()), vgg_type=vgg_type,
|
||||
use_input_norm=use_input_norm)
|
||||
norm_image_with_imagenet_param=norm_image_with_imagenet_param)
|
||||
|
||||
self.percep_criterion, self.style_criterion = self.set_criterion(criterion)
|
||||
|
||||
|
||||
20
loss/gan.py
20
loss/gan.py
@@ -1,4 +1,5 @@
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
@@ -10,7 +11,7 @@ class GANLoss(nn.Module):
|
||||
self.fake_label_val = fake_label_val
|
||||
self.loss_type = loss_type
|
||||
|
||||
def forward(self, prediction, target_is_real: bool, is_discriminator=False):
|
||||
def single_forward(self, prediction, target_is_real: bool, is_discriminator=False):
|
||||
"""
|
||||
gan loss forward
|
||||
:param prediction: network prediction
|
||||
@@ -37,3 +38,20 @@ class GANLoss(nn.Module):
|
||||
return loss
|
||||
else:
|
||||
raise NotImplementedError(f'GAN type {self.loss_type} is not implemented.')
|
||||
|
||||
def forward(self, prediction, target_is_real: bool, is_discriminator=False):
|
||||
if isinstance(prediction, torch.Tensor):
|
||||
# origin
|
||||
return self.single_forward(prediction, target_is_real, is_discriminator)
|
||||
elif isinstance(prediction, list):
|
||||
# for multi scale discriminator, e.g. MultiScaleDiscriminator
|
||||
loss = 0
|
||||
for p in prediction:
|
||||
loss += self.single_forward(p[-1], target_is_real, is_discriminator)
|
||||
return loss
|
||||
elif isinstance(prediction, tuple):
|
||||
# for single discriminator set `need_intermediate_feature` true
|
||||
return self.single_forward(prediction[-1], target_is_real, is_discriminator)
|
||||
else:
|
||||
raise NotImplementedError(f"not support discriminator output: {prediction}")
|
||||
|
||||
|
||||
@@ -1,62 +0,0 @@
|
||||
import torch.nn as nn
|
||||
|
||||
from model.normalization import select_norm_layer
|
||||
from model.registry import MODEL
|
||||
from .base import ResidualBlock
|
||||
|
||||
|
||||
@MODEL.register_module("CyCle-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, base_channels=64, num_blocks=9, padding_mode='reflect',
|
||||
norm_type="IN"):
|
||||
super(Generator, self).__init__()
|
||||
assert num_blocks >= 0, f'Number of residual blocks must be non-negative, but got {num_blocks}.'
|
||||
norm_layer = select_norm_layer(norm_type)
|
||||
use_bias = norm_type == "IN"
|
||||
|
||||
self.start_conv = nn.Sequential(
|
||||
nn.Conv2d(in_channels, base_channels, kernel_size=7, stride=1, padding_mode=padding_mode, padding=3,
|
||||
bias=use_bias),
|
||||
norm_layer(num_features=base_channels),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
|
||||
# down sampling
|
||||
submodules = []
|
||||
num_down_sampling = 2
|
||||
for i in range(num_down_sampling):
|
||||
multiple = 2 ** i
|
||||
submodules += [
|
||||
nn.Conv2d(in_channels=base_channels * multiple, out_channels=base_channels * multiple * 2,
|
||||
kernel_size=3, stride=2, padding=1, bias=use_bias),
|
||||
norm_layer(num_features=base_channels * multiple * 2),
|
||||
nn.ReLU(inplace=True)
|
||||
]
|
||||
self.encoder = nn.Sequential(*submodules)
|
||||
|
||||
res_block_channels = num_down_sampling ** 2 * base_channels
|
||||
self.resnet_middle = nn.Sequential(
|
||||
*[ResidualBlock(res_block_channels, padding_mode, norm_type) for _ in
|
||||
range(num_blocks)])
|
||||
|
||||
# up sampling
|
||||
submodules = []
|
||||
for i in range(num_down_sampling):
|
||||
multiple = 2 ** (num_down_sampling - i)
|
||||
submodules += [
|
||||
nn.ConvTranspose2d(base_channels * multiple, base_channels * multiple // 2, kernel_size=3, stride=2,
|
||||
padding=1, output_padding=1, bias=use_bias),
|
||||
norm_layer(num_features=base_channels * multiple // 2),
|
||||
nn.ReLU(inplace=True),
|
||||
]
|
||||
self.decoder = nn.Sequential(*submodules)
|
||||
|
||||
self.end_conv = nn.Sequential(
|
||||
nn.Conv2d(base_channels, out_channels, kernel_size=7, padding=3, padding_mode=padding_mode),
|
||||
nn.Tanh()
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.encoder(self.start_conv(x))
|
||||
x = self.resnet_middle(x)
|
||||
return self.end_conv(self.decoder(x))
|
||||
@@ -1,150 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from model import MODEL
|
||||
from model.GAN.base import Conv2dBlock, ResBlock
|
||||
from model.normalization import select_norm_layer
|
||||
|
||||
|
||||
class StyleEncoder(nn.Module):
|
||||
def __init__(self, in_channels, out_dim, num_conv, base_channels=64, use_spectral_norm=False,
|
||||
max_multiple=2, padding_mode='reflect', activation_type="ReLU", norm_type="NONE"):
|
||||
super(StyleEncoder, self).__init__()
|
||||
|
||||
sequence = [Conv2dBlock(
|
||||
in_channels, base_channels, kernel_size=7, stride=1, padding=3, padding_mode=padding_mode,
|
||||
use_spectral_norm=use_spectral_norm, activation_type=activation_type, norm_type=norm_type
|
||||
)]
|
||||
|
||||
multiple_now = 1
|
||||
for i in range(1, num_conv + 1):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** max_multiple)
|
||||
sequence.append(Conv2dBlock(
|
||||
multiple_prev * base_channels, multiple_now * base_channels,
|
||||
kernel_size=4, stride=2, padding=1, padding_mode=padding_mode,
|
||||
use_spectral_norm=use_spectral_norm, activation_type=activation_type, norm_type=norm_type
|
||||
))
|
||||
sequence.append(nn.AdaptiveAvgPool2d(1))
|
||||
# conv1x1 works as fc when tensor's size is (batch_size, channels, 1, 1), keep same with origin code
|
||||
sequence.append(nn.Conv2d(multiple_now * base_channels, out_dim, kernel_size=1, stride=1, padding=0))
|
||||
self.model = nn.Sequential(*sequence)
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x).view(x.size(0), -1)
|
||||
|
||||
|
||||
class ContentEncoder(nn.Module):
|
||||
def __init__(self, in_channels, num_down_sampling, num_res_blocks, base_channels=64, use_spectral_norm=False,
|
||||
padding_mode='reflect', activation_type="ReLU", norm_type="IN"):
|
||||
super().__init__()
|
||||
sequence = [Conv2dBlock(
|
||||
in_channels, base_channels, kernel_size=7, stride=1, padding=3, padding_mode=padding_mode,
|
||||
use_spectral_norm=use_spectral_norm, activation_type=activation_type, norm_type=norm_type
|
||||
)]
|
||||
|
||||
for i in range(num_down_sampling):
|
||||
sequence.append(Conv2dBlock(
|
||||
base_channels * (2 ** i), base_channels * (2 ** (i + 1)),
|
||||
kernel_size=4, stride=2, padding=1, padding_mode=padding_mode,
|
||||
use_spectral_norm=use_spectral_norm, activation_type=activation_type, norm_type=norm_type
|
||||
))
|
||||
|
||||
sequence += [ResBlock(base_channels * (2 ** num_down_sampling), use_spectral_norm, padding_mode, norm_type,
|
||||
activation_type) for _ in range(num_res_blocks)]
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
|
||||
def forward(self, x):
|
||||
return self.sequence(x)
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, num_up_sampling, num_res_blocks,
|
||||
use_spectral_norm=False, res_norm_type="AdaIN", norm_type="LN", activation_type="ReLU",
|
||||
padding_mode='reflect'):
|
||||
super(Decoder, self).__init__()
|
||||
self.res_norm_type = res_norm_type
|
||||
self.res_blocks = nn.ModuleList([
|
||||
ResBlock(in_channels, use_spectral_norm, padding_mode, res_norm_type, activation_type=activation_type)
|
||||
for _ in range(num_res_blocks)
|
||||
])
|
||||
sequence = list()
|
||||
channels = in_channels
|
||||
for i in range(num_up_sampling):
|
||||
sequence.append(nn.Sequential(
|
||||
nn.Upsample(scale_factor=2),
|
||||
Conv2dBlock(channels, channels // 2,
|
||||
kernel_size=5, stride=1, padding=2, padding_mode=padding_mode,
|
||||
use_spectral_norm=use_spectral_norm, activation_type=activation_type, norm_type=norm_type
|
||||
),
|
||||
))
|
||||
channels = channels // 2
|
||||
sequence.append(
|
||||
Conv2dBlock(channels, out_channels, kernel_size=7, stride=1, padding=3, padding_mode="reflect",
|
||||
use_spectral_norm=use_spectral_norm, activation_type="Tanh", norm_type="NONE"))
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
|
||||
def forward(self, x):
|
||||
for blk in self.res_blocks:
|
||||
x = blk(x)
|
||||
return self.sequence(x)
|
||||
|
||||
|
||||
class Fusion(nn.Module):
|
||||
def __init__(self, in_features, out_features, base_features, n_blocks, norm_type="NONE"):
|
||||
super().__init__()
|
||||
norm_layer = select_norm_layer(norm_type)
|
||||
self.start_fc = nn.Sequential(
|
||||
nn.Linear(in_features, base_features),
|
||||
norm_layer(base_features),
|
||||
nn.ReLU(True),
|
||||
)
|
||||
self.fcs = nn.Sequential(*[
|
||||
nn.Sequential(
|
||||
nn.Linear(base_features, base_features),
|
||||
norm_layer(base_features),
|
||||
nn.ReLU(True),
|
||||
) for _ in range(n_blocks - 2)
|
||||
])
|
||||
self.end_fc = nn.Sequential(
|
||||
nn.Linear(base_features, out_features),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.start_fc(x)
|
||||
x = self.fcs(x)
|
||||
return self.end_fc(x)
|
||||
|
||||
|
||||
@MODEL.register_module("MUNIT-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, base_channels, num_sampling, num_style_dim, num_style_conv,
|
||||
num_content_res_blocks, num_decoder_res_blocks, num_fusion_dim, num_fusion_blocks,
|
||||
use_spectral_norm=False, activation_type="ReLU", padding_mode='reflect'):
|
||||
super().__init__()
|
||||
self.num_decoder_res_blocks = num_decoder_res_blocks
|
||||
self.content_encoder = ContentEncoder(in_channels, num_sampling, num_content_res_blocks, base_channels,
|
||||
use_spectral_norm, padding_mode, activation_type, norm_type="IN")
|
||||
self.style_encoder = StyleEncoder(in_channels, num_style_dim, num_style_conv, base_channels, use_spectral_norm,
|
||||
padding_mode, activation_type, norm_type="NONE")
|
||||
content_channels = base_channels * (2 ** 2)
|
||||
self.decoder = Decoder(content_channels, out_channels, num_sampling,
|
||||
num_decoder_res_blocks, use_spectral_norm, "AdaIN", norm_type="LN",
|
||||
activation_type=activation_type, padding_mode=padding_mode)
|
||||
self.fusion = Fusion(num_style_dim, num_decoder_res_blocks * 2 * content_channels * 2,
|
||||
base_features=num_fusion_dim, n_blocks=num_fusion_blocks, norm_type="NONE")
|
||||
|
||||
def encode(self, x):
|
||||
return self.content_encoder(x), self.style_encoder(x)
|
||||
|
||||
def decode(self, content, style):
|
||||
as_param_style = torch.chunk(self.fusion(style), self.num_decoder_res_blocks * 2, dim=1)
|
||||
# set style for decoder
|
||||
for i, blk in enumerate(self.decoder.res_blocks):
|
||||
blk.conv1.normalization.set_style(as_param_style[2 * i])
|
||||
blk.conv2.normalization.set_style(as_param_style[2 * i + 1])
|
||||
return self.decoder(content)
|
||||
|
||||
def forward(self, x):
|
||||
content, style = self.encode(x)
|
||||
return self.decode(content, style)
|
||||
@@ -1,118 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchvision.models import vgg19
|
||||
|
||||
from model.normalization import select_norm_layer
|
||||
from model.registry import MODEL
|
||||
from .MUNIT import ContentEncoder, Fusion, Decoder, StyleEncoder
|
||||
from .base import ResBlock
|
||||
|
||||
|
||||
class VGG19StyleEncoder(nn.Module):
|
||||
def __init__(self, in_channels, base_channels=64, style_dim=512, padding_mode='reflect', norm_type="NONE",
|
||||
vgg19_layers=(0, 5, 10, 19), fix_vgg19=True):
|
||||
super().__init__()
|
||||
self.vgg19_layers = vgg19_layers
|
||||
self.vgg19 = vgg19(pretrained=True).features[:vgg19_layers[-1] + 1]
|
||||
self.vgg19.requires_grad_(not fix_vgg19)
|
||||
|
||||
norm_layer = select_norm_layer(norm_type)
|
||||
|
||||
self.conv0 = nn.Sequential(
|
||||
nn.Conv2d(in_channels, base_channels, kernel_size=7, stride=1, padding=3, padding_mode=padding_mode,
|
||||
bias=True),
|
||||
norm_layer(base_channels),
|
||||
nn.ReLU(True),
|
||||
)
|
||||
self.conv = nn.ModuleList([
|
||||
nn.Sequential(
|
||||
nn.Conv2d(base_channels * (2 ** i), base_channels * (2 ** i), kernel_size=4, stride=2, padding=1,
|
||||
padding_mode=padding_mode, bias=True),
|
||||
norm_layer(base_channels),
|
||||
nn.ReLU(True),
|
||||
) for i in range(1, 4)
|
||||
])
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.conv1x1 = nn.Conv2d(base_channels * (2 ** 4), style_dim, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def fixed_style_features(self, x):
|
||||
features = []
|
||||
for i in range(len(self.vgg19)):
|
||||
x = self.vgg19[i](x)
|
||||
if i in self.vgg19_layers:
|
||||
features.append(x)
|
||||
return features
|
||||
|
||||
def forward(self, x):
|
||||
fsf = self.fixed_style_features(x)
|
||||
x = self.conv0(x)
|
||||
for i, l in enumerate(self.conv):
|
||||
x = l(torch.cat([x, fsf[i]], dim=1))
|
||||
x = self.pool(torch.cat([x, fsf[-1]], dim=1))
|
||||
x = self.conv1x1(x)
|
||||
return x.view(x.size(0), -1)
|
||||
|
||||
|
||||
@MODEL.register_module("TAFG-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, style_in_channels, content_in_channels=3, out_channels=3, use_spectral_norm=False,
|
||||
style_encoder_type="StyleEncoder", num_style_conv=4, style_dim=512, num_adain_blocks=8,
|
||||
num_res_blocks=8, base_channels=64, padding_mode="reflect"):
|
||||
super(Generator, self).__init__()
|
||||
self.num_adain_blocks = num_adain_blocks
|
||||
if style_encoder_type == "StyleEncoder":
|
||||
self.style_encoders = nn.ModuleDict(dict(
|
||||
a=StyleEncoder(style_in_channels, style_dim, num_style_conv, base_channels, use_spectral_norm,
|
||||
max_multiple=4, padding_mode=padding_mode, norm_type="NONE"),
|
||||
b=StyleEncoder(style_in_channels, style_dim, num_style_conv, base_channels, use_spectral_norm,
|
||||
max_multiple=4, padding_mode=padding_mode, norm_type="NONE"),
|
||||
))
|
||||
elif style_encoder_type == "VGG19StyleEncoder":
|
||||
self.style_encoders = nn.ModuleDict(dict(
|
||||
a=VGG19StyleEncoder(style_in_channels, base_channels, style_dim=style_dim, padding_mode=padding_mode,
|
||||
norm_type="NONE"),
|
||||
b=VGG19StyleEncoder(style_in_channels, base_channels, style_dim=style_dim, padding_mode=padding_mode,
|
||||
norm_type="NONE", fix_vgg19=False)
|
||||
))
|
||||
else:
|
||||
raise NotImplemented(f"do not support {style_encoder_type}")
|
||||
resnet_channels = 2 ** 2 * base_channels
|
||||
self.style_converters = nn.ModuleDict(dict(
|
||||
a=Fusion(style_dim, num_adain_blocks * 2 * resnet_channels * 2, base_features=256, n_blocks=3,
|
||||
norm_type="NONE"),
|
||||
b=Fusion(style_dim, num_adain_blocks * 2 * resnet_channels * 2, base_features=256, n_blocks=3,
|
||||
norm_type="NONE"),
|
||||
))
|
||||
self.content_encoders = nn.ModuleDict({
|
||||
"a": ContentEncoder(content_in_channels, 2, num_res_blocks=0, use_spectral_norm=use_spectral_norm),
|
||||
"b": ContentEncoder(1, 2, num_res_blocks=0, use_spectral_norm=use_spectral_norm)
|
||||
})
|
||||
|
||||
self.content_resnet = nn.Sequential(*[
|
||||
ResBlock(resnet_channels, use_spectral_norm, padding_mode, "IN")
|
||||
for _ in range(num_res_blocks)
|
||||
])
|
||||
self.decoders = nn.ModuleDict(dict(
|
||||
a=Decoder(resnet_channels, out_channels, 2,
|
||||
num_adain_blocks, use_spectral_norm, "AdaIN", norm_type="LN", padding_mode=padding_mode),
|
||||
b=Decoder(resnet_channels, out_channels, 2,
|
||||
num_adain_blocks, use_spectral_norm, "AdaIN", norm_type="LN", padding_mode=padding_mode),
|
||||
))
|
||||
|
||||
def encode(self, content_img, style_img, which_content, which_style):
|
||||
content = self.content_resnet(self.content_encoders[which_content](content_img))
|
||||
style = self.style_encoders[which_style](style_img)
|
||||
return content, style
|
||||
|
||||
def decode(self, content, style, which):
|
||||
decoder = self.decoders[which]
|
||||
as_param_style = torch.chunk(self.style_converters[which](style), self.num_adain_blocks * 2, dim=1)
|
||||
# set style for decoder
|
||||
for i, blk in enumerate(decoder.res_blocks):
|
||||
blk.conv1.normalization.set_style(as_param_style[2 * i])
|
||||
blk.conv2.normalization.set_style(as_param_style[2 * i + 1])
|
||||
return decoder(content)
|
||||
|
||||
def forward(self, content_img, style_img, which_content, which_style):
|
||||
content, style = self.encode(content_img, style_img, which_content, which_style)
|
||||
return self.decode(content, style, which_style)
|
||||
@@ -1,192 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from model import MODEL
|
||||
from model.normalization import AdaptiveInstanceNorm2d
|
||||
from model.normalization import select_norm_layer
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, padding_mode='zeros', norm_type="IN", use_bias=None,
|
||||
use_spectral=True):
|
||||
super().__init__()
|
||||
self.padding_mode = padding_mode
|
||||
self.use_bias = use_bias
|
||||
self.use_spectral = use_spectral
|
||||
if use_bias is None:
|
||||
# Only for IN, use bias since it does not have affine parameters.
|
||||
self.use_bias = norm_type == "IN"
|
||||
norm_layer = select_norm_layer(norm_type)
|
||||
self.main = nn.Sequential(
|
||||
self.conv_block(in_channels, in_channels),
|
||||
norm_layer(in_channels),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
self.conv_block(in_channels, out_channels),
|
||||
norm_layer(out_channels),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
)
|
||||
self.skip = nn.Sequential(
|
||||
self.conv_block(in_channels, out_channels, padding=0, kernel_size=1),
|
||||
norm_layer(out_channels),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
)
|
||||
|
||||
def conv_block(self, in_channels, out_channels, kernel_size=3, padding=1):
|
||||
conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding,
|
||||
padding_mode=self.padding_mode, bias=self.use_bias)
|
||||
if self.use_spectral:
|
||||
return nn.utils.spectral_norm(conv)
|
||||
else:
|
||||
return conv
|
||||
|
||||
def forward(self, x):
|
||||
return self.main(x) + self.skip(x)
|
||||
|
||||
|
||||
class Interpolation(nn.Module):
|
||||
def __init__(self, scale_factor=None, mode='nearest', align_corners=None):
|
||||
super(Interpolation, self).__init__()
|
||||
self.scale_factor = scale_factor
|
||||
self.mode = mode
|
||||
self.align_corners = align_corners
|
||||
|
||||
def forward(self, x):
|
||||
return F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners,
|
||||
recompute_scale_factor=False)
|
||||
|
||||
def __repr__(self):
|
||||
return f"DownSampling(scale_factor={self.scale_factor}, mode={self.mode}, align_corners={self.align_corners})"
|
||||
|
||||
|
||||
class FADE(nn.Module):
|
||||
def __init__(self, use_spectral, features_channels, in_channels, affine=False, track_running_stats=True):
|
||||
super().__init__()
|
||||
self.bn = nn.BatchNorm2d(num_features=in_channels, affine=affine, track_running_stats=track_running_stats)
|
||||
self.alpha_conv = conv_block(use_spectral, features_channels, in_channels, kernel_size=3, padding=1,
|
||||
padding_mode="zeros")
|
||||
self.beta_conv = conv_block(use_spectral, features_channels, in_channels, kernel_size=3, padding=1,
|
||||
padding_mode="zeros")
|
||||
|
||||
def forward(self, x, feature):
|
||||
alpha = self.alpha_conv(feature)
|
||||
beta = self.beta_conv(feature)
|
||||
x = self.bn(x)
|
||||
return alpha * x + beta
|
||||
|
||||
|
||||
class FADEResBlock(nn.Module):
|
||||
def __init__(self, use_spectral, features_channels, in_channels, out_channels):
|
||||
super().__init__()
|
||||
self.main = nn.Sequential(
|
||||
FADE(use_spectral, features_channels, in_channels),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
conv_block(use_spectral, in_channels, in_channels, kernel_size=3, padding=1),
|
||||
FADE(use_spectral, features_channels, in_channels),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
conv_block(use_spectral, in_channels, out_channels, kernel_size=3, padding=1),
|
||||
)
|
||||
self.skip = nn.Sequential(
|
||||
FADE(use_spectral, features_channels, in_channels),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
conv_block(use_spectral, in_channels, out_channels, kernel_size=1, padding=0),
|
||||
)
|
||||
self.up_sample = Interpolation(2, mode="nearest")
|
||||
|
||||
@staticmethod
|
||||
def forward_with_fade(module, x, feature):
|
||||
for layer in module:
|
||||
if layer.__class__.__name__ == "FADE":
|
||||
x = layer(x, feature)
|
||||
else:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
def forward(self, x, feature):
|
||||
out = self.forward_with_fade(self.main, x, feature) + self.forward_with_fade(self.main, x, feature)
|
||||
return self.up_sample(out)
|
||||
|
||||
|
||||
def conv_block(use_spectral, in_channels, out_channels, **kwargs):
|
||||
conv = nn.Conv2d(in_channels, out_channels, **kwargs)
|
||||
return nn.utils.spectral_norm(conv) if use_spectral else conv
|
||||
|
||||
|
||||
@MODEL.register_module("TSIT-Generator")
|
||||
class TSITGenerator(nn.Module):
|
||||
def __init__(self, num_blocks=7, base_channels=64, content_in_channels=3, style_in_channels=3,
|
||||
out_channels=3, use_spectral=True, input_layer_type="conv1x1"):
|
||||
super().__init__()
|
||||
self.num_blocks = num_blocks
|
||||
self.base_channels = base_channels
|
||||
self.use_spectral = use_spectral
|
||||
|
||||
self.content_input_layer = self.build_input_layer(content_in_channels, base_channels, input_layer_type)
|
||||
self.style_input_layer = self.build_input_layer(style_in_channels, base_channels, input_layer_type)
|
||||
self.content_stream = self.build_stream()
|
||||
self.style_stream = self.build_stream()
|
||||
self.generator = self.build_generator()
|
||||
self.end_conv = nn.Sequential(
|
||||
conv_block(use_spectral, base_channels, out_channels, kernel_size=7, padding=3, padding_mode="zeros"),
|
||||
nn.Tanh()
|
||||
)
|
||||
|
||||
def build_generator(self):
|
||||
stream_sequence = []
|
||||
multiple_now = min(2 ** self.num_blocks, 2 ** 4)
|
||||
for i in range(1, self.num_blocks + 1):
|
||||
m = self.num_blocks - i
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** m, 2 ** 4)
|
||||
stream_sequence.append(nn.Sequential(
|
||||
AdaptiveInstanceNorm2d(multiple_prev * self.base_channels),
|
||||
FADEResBlock(self.use_spectral, multiple_prev * self.base_channels, multiple_prev * self.base_channels,
|
||||
multiple_now * self.base_channels)
|
||||
))
|
||||
return nn.ModuleList(stream_sequence)
|
||||
|
||||
def build_input_layer(self, in_channels, out_channels, input_layer_type="conv7x7"):
|
||||
if input_layer_type == "conv7x7":
|
||||
return nn.Sequential(
|
||||
conv_block(self.use_spectral, in_channels, out_channels, kernel_size=7, stride=1,
|
||||
padding_mode="zeros", padding=3, bias=True),
|
||||
select_norm_layer("IN")(out_channels),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
elif input_layer_type == "conv1x1":
|
||||
return conv_block(self.use_spectral, in_channels, out_channels, kernel_size=1, stride=1, padding=0)
|
||||
else:
|
||||
raise NotImplemented
|
||||
|
||||
def build_stream(self):
|
||||
multiple_now = 1
|
||||
stream_sequence = []
|
||||
for i in range(1, self.num_blocks + 1):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** 4)
|
||||
stream_sequence.append(nn.Sequential(
|
||||
Interpolation(scale_factor=0.5, mode="nearest"),
|
||||
ResBlock(multiple_prev * self.base_channels, multiple_now * self.base_channels,
|
||||
use_spectral=self.use_spectral)
|
||||
))
|
||||
return nn.ModuleList(stream_sequence)
|
||||
|
||||
def forward(self, content_img, style_img):
|
||||
c = self.content_input_layer(content_img)
|
||||
s = self.style_input_layer(style_img)
|
||||
content_features = []
|
||||
style_features = []
|
||||
for i in range(self.num_blocks):
|
||||
s = self.style_stream[i](s)
|
||||
c = self.content_stream[i](c)
|
||||
content_features.append(c)
|
||||
style_features.append(s)
|
||||
z = torch.randn(size=content_features[-1].size(), device=content_features[-1].device)
|
||||
|
||||
for i in range(self.num_blocks):
|
||||
m = - i - 1
|
||||
layer = self.generator[i]
|
||||
layer[0].set_style(torch.cat(torch.std_mean(style_features[m], dim=[2, 3]), dim=1))
|
||||
z = layer[0](z)
|
||||
z = layer[1](z, content_features[m])
|
||||
return self.end_conv(z)
|
||||
@@ -1,236 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from .base import ResidualBlock
|
||||
from model.registry import MODEL
|
||||
|
||||
|
||||
class RhoClipper(object):
|
||||
def __init__(self, clip_min, clip_max):
|
||||
self.clip_min = clip_min
|
||||
self.clip_max = clip_max
|
||||
assert clip_min < clip_max
|
||||
|
||||
def __call__(self, module):
|
||||
if hasattr(module, 'rho'):
|
||||
w = module.rho.data
|
||||
w = w.clamp(self.clip_min, self.clip_max)
|
||||
module.rho.data = w
|
||||
|
||||
|
||||
@MODEL.register_module("UGATIT-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, base_channels=64, num_blocks=6, img_size=256, light=False):
|
||||
assert (num_blocks >= 0)
|
||||
super(Generator, self).__init__()
|
||||
self.input_channels = in_channels
|
||||
self.output_channels = out_channels
|
||||
self.base_channels = base_channels
|
||||
self.num_blocks = num_blocks
|
||||
self.img_size = img_size
|
||||
self.light = light
|
||||
|
||||
down_encoder = [nn.Conv2d(in_channels, base_channels, kernel_size=7, stride=1, padding=3,
|
||||
padding_mode="reflect", bias=False),
|
||||
nn.InstanceNorm2d(base_channels),
|
||||
nn.ReLU(True)]
|
||||
|
||||
n_down_sampling = 2
|
||||
for i in range(n_down_sampling):
|
||||
mult = 2 ** i
|
||||
down_encoder += [nn.Conv2d(base_channels * mult, base_channels * mult * 2, kernel_size=3, stride=2,
|
||||
padding=1, bias=False, padding_mode="reflect"),
|
||||
nn.InstanceNorm2d(base_channels * mult * 2),
|
||||
nn.ReLU(True)]
|
||||
|
||||
# Down-Sampling Bottleneck
|
||||
mult = 2 ** n_down_sampling
|
||||
for i in range(num_blocks):
|
||||
down_encoder += [ResidualBlock(base_channels * mult, use_bias=False)]
|
||||
self.down_encoder = nn.Sequential(*down_encoder)
|
||||
|
||||
# Class Activation Map
|
||||
self.gap_fc = nn.Linear(base_channels * mult, 1, bias=False)
|
||||
self.gmp_fc = nn.Linear(base_channels * mult, 1, bias=False)
|
||||
self.conv1x1 = nn.Conv2d(base_channels * mult * 2, base_channels * mult, kernel_size=1, stride=1, bias=True)
|
||||
self.relu = nn.ReLU(True)
|
||||
|
||||
# Gamma, Beta block
|
||||
if self.light:
|
||||
fc = [nn.Linear(base_channels * mult, base_channels * mult, bias=False),
|
||||
nn.ReLU(True),
|
||||
nn.Linear(base_channels * mult, base_channels * mult, bias=False),
|
||||
nn.ReLU(True)]
|
||||
else:
|
||||
fc = [
|
||||
nn.Linear(img_size // mult * img_size // mult * base_channels * mult, base_channels * mult, bias=False),
|
||||
nn.ReLU(True),
|
||||
nn.Linear(base_channels * mult, base_channels * mult, bias=False),
|
||||
nn.ReLU(True)]
|
||||
self.fc = nn.Sequential(*fc)
|
||||
|
||||
self.gamma = nn.Linear(base_channels * mult, base_channels * mult, bias=False)
|
||||
self.beta = nn.Linear(base_channels * mult, base_channels * mult, bias=False)
|
||||
|
||||
# Up-Sampling Bottleneck
|
||||
self.up_bottleneck = nn.ModuleList(
|
||||
[ResnetAdaILNBlock(base_channels * mult, use_bias=False) for _ in range(num_blocks)])
|
||||
|
||||
# Up-Sampling
|
||||
up_decoder = []
|
||||
for i in range(n_down_sampling):
|
||||
mult = 2 ** (n_down_sampling - i)
|
||||
up_decoder += [nn.Upsample(scale_factor=2, mode='nearest'),
|
||||
nn.Conv2d(base_channels * mult, base_channels * mult // 2, kernel_size=3, stride=1,
|
||||
padding=1, padding_mode="reflect", bias=False),
|
||||
ILN(base_channels * mult // 2),
|
||||
nn.ReLU(True)]
|
||||
|
||||
up_decoder += [nn.Conv2d(base_channels, out_channels, kernel_size=7, stride=1, padding=3,
|
||||
padding_mode="reflect", bias=False),
|
||||
nn.Tanh()]
|
||||
self.up_decoder = nn.Sequential(*up_decoder)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.down_encoder(x)
|
||||
|
||||
gap = torch.nn.functional.adaptive_avg_pool2d(x, 1)
|
||||
gap_logit = self.gap_fc(gap.view(x.shape[0], -1))
|
||||
gap = x * self.gap_fc.weight.unsqueeze(2).unsqueeze(3)
|
||||
|
||||
gmp = torch.nn.functional.adaptive_max_pool2d(x, 1)
|
||||
gmp_logit = self.gmp_fc(gmp.view(x.shape[0], -1))
|
||||
gmp = x * self.gmp_fc.weight.unsqueeze(2).unsqueeze(3)
|
||||
|
||||
cam_logit = torch.cat([gap_logit, gmp_logit], 1)
|
||||
|
||||
x = torch.cat([gap, gmp], 1)
|
||||
x = self.relu(self.conv1x1(x))
|
||||
|
||||
heatmap = torch.sum(x, dim=1, keepdim=True)
|
||||
|
||||
if self.light:
|
||||
x_ = torch.nn.functional.adaptive_avg_pool2d(x, 1)
|
||||
x_ = self.fc(x_.view(x_.shape[0], -1))
|
||||
else:
|
||||
x_ = self.fc(x.view(x.shape[0], -1))
|
||||
gamma, beta = self.gamma(x_), self.beta(x_)
|
||||
|
||||
for ub in self.up_bottleneck:
|
||||
x = ub(x, gamma, beta)
|
||||
|
||||
x = self.up_decoder(x)
|
||||
return x, cam_logit, heatmap
|
||||
|
||||
|
||||
class ResnetAdaILNBlock(nn.Module):
|
||||
def __init__(self, dim, use_bias):
|
||||
super(ResnetAdaILNBlock, self).__init__()
|
||||
self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, bias=use_bias, padding_mode="reflect")
|
||||
self.norm1 = AdaILN(dim)
|
||||
self.relu1 = nn.ReLU(True)
|
||||
|
||||
self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, stride=1, padding=1, bias=use_bias, padding_mode="reflect")
|
||||
self.norm2 = AdaILN(dim)
|
||||
|
||||
def forward(self, x, gamma, beta):
|
||||
out = self.conv1(x)
|
||||
out = self.norm1(out, gamma, beta)
|
||||
out = self.relu1(out)
|
||||
out = self.conv2(out)
|
||||
out = self.norm2(out, gamma, beta)
|
||||
|
||||
return out + x
|
||||
|
||||
|
||||
def instance_layer_normalization(x, gamma, beta, rho, eps=1e-5):
|
||||
in_mean, in_var = torch.mean(x, dim=[2, 3], keepdim=True), torch.var(x, dim=[2, 3], keepdim=True)
|
||||
out_in = (x - in_mean) / torch.sqrt(in_var + eps)
|
||||
ln_mean, ln_var = torch.mean(x, dim=[1, 2, 3], keepdim=True), torch.var(x, dim=[1, 2, 3], keepdim=True)
|
||||
out_ln = (x - ln_mean) / torch.sqrt(ln_var + eps)
|
||||
out = rho.expand(x.shape[0], -1, -1, -1) * out_in + (1 - rho.expand(x.shape[0], -1, -1, -1)) * out_ln
|
||||
out = out * gamma.unsqueeze(2).unsqueeze(3) + beta.unsqueeze(2).unsqueeze(3)
|
||||
return out
|
||||
|
||||
|
||||
class AdaILN(nn.Module):
|
||||
def __init__(self, num_features, eps=1e-5, default_rho=0.9):
|
||||
super(AdaILN, self).__init__()
|
||||
self.eps = eps
|
||||
self.rho = nn.Parameter(torch.Tensor(1, num_features, 1, 1))
|
||||
self.rho.data.fill_(default_rho)
|
||||
|
||||
def forward(self, x, gamma, beta):
|
||||
return instance_layer_normalization(x, gamma, beta, self.rho, self.eps)
|
||||
|
||||
|
||||
class ILN(nn.Module):
|
||||
def __init__(self, num_features, eps=1e-5):
|
||||
super(ILN, self).__init__()
|
||||
self.eps = eps
|
||||
self.rho = nn.Parameter(torch.Tensor(1, num_features, 1, 1))
|
||||
self.gamma = nn.Parameter(torch.Tensor(1, num_features))
|
||||
self.beta = nn.Parameter(torch.Tensor(1, num_features))
|
||||
self.rho.data.fill_(0.0)
|
||||
self.gamma.data.fill_(1.0)
|
||||
self.beta.data.fill_(0.0)
|
||||
|
||||
def forward(self, x):
|
||||
return instance_layer_normalization(
|
||||
x, self.gamma.expand(x.shape[0], -1), self.beta.expand(x.shape[0], -1), self.rho, self.eps)
|
||||
|
||||
|
||||
@MODEL.register_module("UGATIT-Discriminator")
|
||||
class Discriminator(nn.Module):
|
||||
def __init__(self, in_channels, base_channels=64, num_blocks=5):
|
||||
super(Discriminator, self).__init__()
|
||||
encoder = [self.build_conv_block(in_channels, base_channels)]
|
||||
|
||||
for i in range(1, num_blocks - 2):
|
||||
mult = 2 ** (i - 1)
|
||||
encoder.append(self.build_conv_block(base_channels * mult, base_channels * mult * 2))
|
||||
|
||||
mult = 2 ** (num_blocks - 2 - 1)
|
||||
encoder.append(self.build_conv_block(base_channels * mult, base_channels * mult * 2, stride=1))
|
||||
|
||||
self.encoder = nn.Sequential(*encoder)
|
||||
|
||||
# Class Activation Map
|
||||
mult = 2 ** (num_blocks - 2)
|
||||
self.gap_fc = nn.utils.spectral_norm(nn.Linear(base_channels * mult, 1, bias=False))
|
||||
self.gmp_fc = nn.utils.spectral_norm(nn.Linear(base_channels * mult, 1, bias=False))
|
||||
self.conv1x1 = nn.Conv2d(base_channels * mult * 2, base_channels * mult, kernel_size=1, stride=1, bias=True)
|
||||
self.leaky_relu = nn.LeakyReLU(0.2, True)
|
||||
|
||||
self.conv = nn.utils.spectral_norm(
|
||||
nn.Conv2d(base_channels * mult, 1, kernel_size=4, stride=1, padding=1, bias=False, padding_mode="reflect"))
|
||||
|
||||
@staticmethod
|
||||
def build_conv_block(in_channels, out_channels, kernel_size=4, stride=2, padding=1, padding_mode="reflect"):
|
||||
return nn.Sequential(*[
|
||||
nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride,
|
||||
bias=True, padding=padding, padding_mode=padding_mode)),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
])
|
||||
|
||||
def forward(self, x, return_heatmap=False):
|
||||
x = self.encoder(x)
|
||||
batch_size = x.size(0)
|
||||
|
||||
gap = torch.nn.functional.adaptive_avg_pool2d(x, 1) # B x C x 1 x 1, avg of per channel
|
||||
gap_logit = self.gap_fc(gap.view(batch_size, -1))
|
||||
gap = x * self.gap_fc.weight.unsqueeze(2).unsqueeze(3)
|
||||
|
||||
gmp = torch.nn.functional.adaptive_max_pool2d(x, 1)
|
||||
gmp_logit = self.gmp_fc(gmp.view(batch_size, -1))
|
||||
gmp = x * self.gmp_fc.weight.unsqueeze(2).unsqueeze(3)
|
||||
|
||||
cam_logit = torch.cat([gap_logit, gmp_logit], 1)
|
||||
|
||||
x = torch.cat([gap, gmp], 1)
|
||||
x = self.leaky_relu(self.conv1x1(x))
|
||||
|
||||
if return_heatmap:
|
||||
heatmap = torch.sum(x, dim=1, keepdim=True)
|
||||
return self.conv(x), cam_logit, heatmap
|
||||
else:
|
||||
return self.conv(x), cam_logit
|
||||
@@ -1,203 +0,0 @@
|
||||
from functools import partial
|
||||
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from model import MODEL
|
||||
from model.normalization import select_norm_layer
|
||||
|
||||
|
||||
class GANImageBuffer(object):
|
||||
"""This class implements an image buffer that stores previously
|
||||
generated images.
|
||||
This buffer allows us to update the discriminator using a history of
|
||||
generated images rather than the ones produced by the latest generator
|
||||
to reduce model oscillation.
|
||||
Args:
|
||||
buffer_size (int): The size of image buffer. If buffer_size = 0,
|
||||
no buffer will be created.
|
||||
buffer_ratio (float): The chance / possibility to use the images
|
||||
previously stored in the buffer.
|
||||
"""
|
||||
|
||||
def __init__(self, buffer_size, buffer_ratio=0.5):
|
||||
self.buffer_size = buffer_size
|
||||
# create an empty buffer
|
||||
if self.buffer_size > 0:
|
||||
self.img_num = 0
|
||||
self.image_buffer = []
|
||||
self.buffer_ratio = buffer_ratio
|
||||
|
||||
def query(self, images):
|
||||
"""Query current image batch using a history of generated images.
|
||||
Args:
|
||||
images (Tensor): Current image batch without history information.
|
||||
"""
|
||||
if self.buffer_size == 0: # if the buffer size is 0, do nothing
|
||||
return images
|
||||
return_images = []
|
||||
for image in images:
|
||||
image = torch.unsqueeze(image.data, 0)
|
||||
# if the buffer is not full, keep inserting current images
|
||||
if self.img_num < self.buffer_size:
|
||||
self.img_num = self.img_num + 1
|
||||
self.image_buffer.append(image)
|
||||
return_images.append(image)
|
||||
else:
|
||||
use_buffer = torch.rand(1) < self.buffer_ratio
|
||||
# by self.buffer_ratio, the buffer will return a previously
|
||||
# stored image, and insert the current image into the buffer
|
||||
if use_buffer:
|
||||
random_id = torch.randint(0, self.buffer_size, (1,)).item()
|
||||
image_tmp = self.image_buffer[random_id].clone()
|
||||
self.image_buffer[random_id] = image
|
||||
return_images.append(image_tmp)
|
||||
# by (1 - self.buffer_ratio), the buffer will return the
|
||||
# current image
|
||||
else:
|
||||
return_images.append(image)
|
||||
# collect all the images and return
|
||||
return_images = torch.cat(return_images, 0)
|
||||
return return_images
|
||||
|
||||
|
||||
# based SPADE or pix2pixHD Discriminator
|
||||
@MODEL.register_module("PatchDiscriminator")
|
||||
class PatchDiscriminator(nn.Module):
|
||||
def __init__(self, in_channels, base_channels, num_conv=4, use_spectral=False, norm_type="IN",
|
||||
need_intermediate_feature=False):
|
||||
super().__init__()
|
||||
self.need_intermediate_feature = need_intermediate_feature
|
||||
|
||||
kernel_size = 4
|
||||
padding = math.ceil((kernel_size - 1.0) / 2)
|
||||
norm_layer = select_norm_layer(norm_type)
|
||||
use_bias = norm_type == "IN"
|
||||
padding_mode = "zeros"
|
||||
|
||||
sequence = [nn.Sequential(
|
||||
nn.Conv2d(in_channels, base_channels, kernel_size, stride=2, padding=padding),
|
||||
nn.LeakyReLU(0.2, False)
|
||||
)]
|
||||
multiple_now = 1
|
||||
for i in range(1, num_conv):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** 3)
|
||||
stride = 1 if i == num_conv - 1 else 2
|
||||
sequence.append(nn.Sequential(
|
||||
self.build_conv2d(use_spectral, base_channels * multiple_prev, base_channels * multiple_now,
|
||||
kernel_size, stride, padding, bias=use_bias, padding_mode=padding_mode),
|
||||
norm_layer(base_channels * multiple_now),
|
||||
nn.LeakyReLU(0.2, inplace=False),
|
||||
))
|
||||
multiple_now = min(2 ** num_conv, 8)
|
||||
sequence.append(nn.Conv2d(base_channels * multiple_now, 1, kernel_size, stride=1, padding=padding,
|
||||
padding_mode=padding_mode))
|
||||
self.conv_blocks = nn.ModuleList(sequence)
|
||||
|
||||
@staticmethod
|
||||
def build_conv2d(use_spectral, in_channels: int, out_channels: int, kernel_size, stride, padding,
|
||||
bias=True, padding_mode: str = 'zeros'):
|
||||
conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias, padding_mode=padding_mode)
|
||||
if not use_spectral:
|
||||
return conv
|
||||
return nn.utils.spectral_norm(conv)
|
||||
|
||||
def forward(self, x):
|
||||
if self.need_intermediate_feature:
|
||||
intermediate_feature = []
|
||||
for layer in self.conv_blocks:
|
||||
x = layer(x)
|
||||
intermediate_feature.append(x)
|
||||
return tuple(intermediate_feature)
|
||||
else:
|
||||
for layer in self.conv_blocks:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
@MODEL.register_module()
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, num_channels, padding_mode='reflect', norm_type="IN", use_bias=None):
|
||||
super(ResidualBlock, self).__init__()
|
||||
if use_bias is None:
|
||||
# Only for IN, use bias since it does not have affine parameters.
|
||||
use_bias = norm_type == "IN"
|
||||
norm_layer = select_norm_layer(norm_type)
|
||||
self.conv1 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode,
|
||||
bias=use_bias)
|
||||
self.norm1 = norm_layer(num_channels)
|
||||
self.relu1 = nn.ReLU(inplace=True)
|
||||
self.conv2 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, padding_mode=padding_mode,
|
||||
bias=use_bias)
|
||||
self.norm2 = norm_layer(num_channels)
|
||||
|
||||
def forward(self, x):
|
||||
res = x
|
||||
x = self.relu1(self.norm1(self.conv1(x)))
|
||||
x = self.norm2(self.conv2(x))
|
||||
return x + res
|
||||
|
||||
|
||||
_DO_NO_THING_FUNC = lambda x: x
|
||||
|
||||
|
||||
def select_activation(t):
|
||||
if t == "ReLU":
|
||||
return partial(nn.ReLU, inplace=True)
|
||||
elif t == "LeakyReLU":
|
||||
return partial(nn.LeakyReLU, negative_slope=0.2, inplace=True)
|
||||
elif t == "Tanh":
|
||||
return partial(nn.Tanh)
|
||||
elif t == "NONE":
|
||||
return _DO_NO_THING_FUNC
|
||||
else:
|
||||
raise NotImplemented
|
||||
|
||||
|
||||
def _use_bias_checker(norm_type):
|
||||
return norm_type not in ["IN", "BN", "AdaIN"]
|
||||
|
||||
|
||||
class Conv2dBlock(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int, use_spectral_norm=False, activation_type="ReLU",
|
||||
bias=None, norm_type="NONE", **conv_kwargs):
|
||||
super().__init__()
|
||||
self.norm_type = norm_type
|
||||
self.activation_type = activation_type
|
||||
conv_kwargs["bias"] = _use_bias_checker(norm_type) if bias is None else bias
|
||||
conv = nn.Conv2d(in_channels, out_channels, **conv_kwargs)
|
||||
self.convolution = nn.utils.spectral_norm(conv) if use_spectral_norm else conv
|
||||
if norm_type != "NONE":
|
||||
self.normalization = select_norm_layer(norm_type)(out_channels)
|
||||
if activation_type != "NONE":
|
||||
self.activation = select_activation(activation_type)()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.convolution(x)
|
||||
if self.norm_type != "NONE":
|
||||
x = self.normalization(x)
|
||||
if self.activation_type != "NONE":
|
||||
x = self.activation(x)
|
||||
return x
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, num_channels, use_spectral_norm=False, padding_mode='reflect',
|
||||
norm_type="IN", activation_type="ReLU", use_bias=None):
|
||||
super().__init__()
|
||||
self.norm_type = norm_type
|
||||
if use_bias is None:
|
||||
# bias will be canceled after channel wise normalization
|
||||
use_bias = _use_bias_checker(norm_type)
|
||||
|
||||
self.conv1 = Conv2dBlock(num_channels, num_channels, use_spectral_norm,
|
||||
kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias,
|
||||
norm_type=norm_type, activation_type=activation_type)
|
||||
self.conv2 = Conv2dBlock(num_channels, num_channels, use_spectral_norm,
|
||||
kernel_size=3, padding=1, padding_mode=padding_mode, bias=use_bias,
|
||||
norm_type=norm_type, activation_type="NONE")
|
||||
|
||||
def forward(self, x):
|
||||
return self.conv2(self.conv1(x)) + x
|
||||
@@ -1,8 +1,7 @@
|
||||
from model.registry import MODEL
|
||||
import model.GAN.CycleGAN
|
||||
import model.GAN.TAFG
|
||||
import model.GAN.UGATIT
|
||||
import model.GAN.wrapper
|
||||
import model.GAN.base
|
||||
import model.GAN.TSIT
|
||||
import model.GAN.MUNIT
|
||||
from model.registry import MODEL, NORMALIZATION
|
||||
import model.base.normalization
|
||||
import model.image_translation.UGATIT
|
||||
import model.image_translation.CycleGAN
|
||||
import model.image_translation.pix2pixHD
|
||||
import model.image_translation.GauGAN
|
||||
import model.image_translation.TSIT
|
||||
128
model/base/module.py
Normal file
128
model/base/module.py
Normal file
@@ -0,0 +1,128 @@
|
||||
import torch.nn as nn
|
||||
|
||||
from model.registry import NORMALIZATION
|
||||
|
||||
_DO_NO_THING_FUNC = lambda x: x
|
||||
|
||||
|
||||
def _use_bias_checker(norm_type):
|
||||
return norm_type not in ["IN", "BN", "AdaIN", "FADE", "SPADE"]
|
||||
|
||||
|
||||
def _normalization(norm, num_features, additional_kwargs=None):
|
||||
if norm == "NONE":
|
||||
return _DO_NO_THING_FUNC
|
||||
|
||||
if additional_kwargs is None:
|
||||
additional_kwargs = {}
|
||||
kwargs = dict(_type=norm, num_features=num_features)
|
||||
kwargs.update(additional_kwargs)
|
||||
return NORMALIZATION.build_with(kwargs)
|
||||
|
||||
|
||||
def _activation(activation, inplace=True):
|
||||
if activation == "NONE":
|
||||
return _DO_NO_THING_FUNC
|
||||
elif activation == "ReLU":
|
||||
return nn.ReLU(inplace=inplace)
|
||||
elif activation == "LeakyReLU":
|
||||
return nn.LeakyReLU(negative_slope=0.2, inplace=inplace)
|
||||
elif activation == "Tanh":
|
||||
return nn.Tanh()
|
||||
else:
|
||||
raise NotImplementedError(f"{activation} not valid")
|
||||
|
||||
|
||||
class LinearBlock(nn.Module):
|
||||
def __init__(self, in_features: int, out_features: int, bias=None, activation_type="ReLU", norm_type="NONE"):
|
||||
super().__init__()
|
||||
|
||||
self.norm_type = norm_type
|
||||
self.activation_type = activation_type
|
||||
|
||||
bias = _use_bias_checker(norm_type) if bias is None else bias
|
||||
self.linear = nn.Linear(in_features, out_features, bias)
|
||||
|
||||
self.normalization = _normalization(norm_type, out_features)
|
||||
self.activation = _activation(activation_type)
|
||||
|
||||
def forward(self, x):
|
||||
return self.activation(self.normalization(self.linear(x)))
|
||||
|
||||
|
||||
class Conv2dBlock(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int, bias=None,
|
||||
activation_type="ReLU", norm_type="NONE", additional_norm_kwargs=None,
|
||||
pre_activation=False, use_transpose_conv=False, **conv_kwargs):
|
||||
super().__init__()
|
||||
self.norm_type = norm_type
|
||||
self.activation_type = activation_type
|
||||
self.pre_activation = pre_activation
|
||||
|
||||
if use_transpose_conv:
|
||||
# Only "zeros" padding mode is supported for ConvTranspose2d
|
||||
conv_kwargs["padding_mode"] = "zeros"
|
||||
conv = nn.ConvTranspose2d
|
||||
else:
|
||||
conv = nn.Conv2d
|
||||
|
||||
if pre_activation:
|
||||
self.normalization = _normalization(norm_type, in_channels, additional_norm_kwargs)
|
||||
self.activation = _activation(activation_type, inplace=False)
|
||||
self.convolution = conv(in_channels, out_channels, **conv_kwargs)
|
||||
else:
|
||||
# if caller not set bias, set bias automatically.
|
||||
conv_kwargs["bias"] = _use_bias_checker(norm_type) if bias is None else bias
|
||||
self.convolution = conv(in_channels, out_channels, **conv_kwargs)
|
||||
self.normalization = _normalization(norm_type, out_channels, additional_norm_kwargs)
|
||||
self.activation = _activation(activation_type)
|
||||
|
||||
def forward(self, x):
|
||||
if self.pre_activation:
|
||||
return self.convolution(self.activation(self.normalization(x)))
|
||||
return self.activation(self.normalization(self.convolution(x)))
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_channels,
|
||||
padding_mode='reflect', activation_type="ReLU", norm_type="IN", pre_activation=False,
|
||||
out_channels=None, out_activation_type=None, additional_norm_kwargs=None):
|
||||
"""
|
||||
Residual Conv Block
|
||||
:param in_channels:
|
||||
:param out_channels:
|
||||
:param padding_mode:
|
||||
:param activation_type:
|
||||
:param norm_type:
|
||||
:param out_activation_type:
|
||||
:param pre_activation: full pre-activation mode from https://arxiv.org/pdf/1603.05027v3.pdf, figure 4
|
||||
"""
|
||||
super().__init__()
|
||||
self.norm_type = norm_type
|
||||
|
||||
if out_channels is None:
|
||||
out_channels = in_channels
|
||||
if out_activation_type is None:
|
||||
# if not specify `out_activation_type`, using default `out_activation_type`
|
||||
# `out_activation_type` default mode:
|
||||
# "NONE" for not full pre-activation
|
||||
# `norm_type` for full pre-activation
|
||||
out_activation_type = "NONE" if not pre_activation else norm_type
|
||||
|
||||
self.learn_skip_connection = in_channels != out_channels
|
||||
|
||||
conv_param = dict(kernel_size=3, padding=1, norm_type=norm_type, activation_type=activation_type,
|
||||
additional_norm_kwargs=additional_norm_kwargs, pre_activation=pre_activation,
|
||||
padding_mode=padding_mode)
|
||||
|
||||
self.conv1 = Conv2dBlock(in_channels, in_channels, **conv_param)
|
||||
self.conv2 = Conv2dBlock(in_channels, out_channels, **conv_param)
|
||||
|
||||
if self.learn_skip_connection:
|
||||
conv_param['kernel_size'] = 1
|
||||
conv_param['padding'] = 0
|
||||
self.res_conv = Conv2dBlock(in_channels, out_channels, **conv_param)
|
||||
|
||||
def forward(self, x):
|
||||
res = x if not self.learn_skip_connection else self.res_conv(x)
|
||||
return self.conv2(self.conv1(x)) + res
|
||||
143
model/base/normalization.py
Normal file
143
model/base/normalization.py
Normal file
@@ -0,0 +1,143 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from model import NORMALIZATION
|
||||
from model.base.module import Conv2dBlock
|
||||
|
||||
_VALID_NORM_AND_ABBREVIATION = dict(
|
||||
IN="InstanceNorm2d",
|
||||
BN="BatchNorm2d",
|
||||
)
|
||||
|
||||
for abbr, name in _VALID_NORM_AND_ABBREVIATION.items():
|
||||
NORMALIZATION.register_module(module=getattr(nn, name), name=abbr)
|
||||
|
||||
|
||||
@NORMALIZATION.register_module("ADE")
|
||||
class AdaptiveDenormalization(nn.Module):
|
||||
def __init__(self, num_features, base_norm_type="BN", gamma_bias=0.0):
|
||||
super().__init__()
|
||||
self.num_features = num_features
|
||||
self.base_norm_type = base_norm_type
|
||||
self.norm = self.base_norm(num_features)
|
||||
self.gamma = None
|
||||
self.gamma_bias = gamma_bias
|
||||
self.beta = None
|
||||
self.have_set_condition = False
|
||||
|
||||
def base_norm(self, num_features):
|
||||
if self.base_norm_type == "IN":
|
||||
return nn.InstanceNorm2d(num_features, affine=False)
|
||||
elif self.base_norm_type == "BN":
|
||||
return nn.BatchNorm2d(num_features, affine=False, track_running_stats=True)
|
||||
|
||||
def set_condition(self, gamma, beta):
|
||||
self.gamma, self.beta = gamma, beta
|
||||
self.have_set_condition = True
|
||||
|
||||
def forward(self, x):
|
||||
assert self.have_set_condition
|
||||
x = self.norm(x)
|
||||
x = (self.gamma + self.gamma_bias) * x + self.beta
|
||||
self.have_set_condition = False
|
||||
return x
|
||||
#
|
||||
# def __repr__(self):
|
||||
# return f"{self.__class__.__name__}(num_features={self.num_features}, " \
|
||||
# f"base_norm_type={self.base_norm_type})"
|
||||
|
||||
|
||||
@NORMALIZATION.register_module("AdaIN")
|
||||
class AdaptiveInstanceNorm2d(AdaptiveDenormalization):
|
||||
def __init__(self, num_features: int):
|
||||
super().__init__(num_features, "IN")
|
||||
self.num_features = num_features
|
||||
|
||||
def set_style(self, style):
|
||||
style = style.view(*style.size(), 1, 1)
|
||||
gamma, beta = style.chunk(2, 1)
|
||||
super().set_condition(gamma, beta)
|
||||
|
||||
|
||||
@NORMALIZATION.register_module("FADE")
|
||||
class FeatureAdaptiveDenormalization(AdaptiveDenormalization):
|
||||
def __init__(self, num_features: int, condition_in_channels,
|
||||
base_norm_type="BN", padding_mode="zeros", gamma_bias=0.0):
|
||||
super().__init__(num_features, base_norm_type, gamma_bias=gamma_bias)
|
||||
self.beta_conv = nn.Conv2d(condition_in_channels, self.num_features, kernel_size=3, padding=1,
|
||||
padding_mode=padding_mode)
|
||||
self.gamma_conv = nn.Conv2d(condition_in_channels, self.num_features, kernel_size=3, padding=1,
|
||||
padding_mode=padding_mode)
|
||||
|
||||
def set_feature(self, feature):
|
||||
gamma = self.gamma_conv(feature)
|
||||
beta = self.beta_conv(feature)
|
||||
super().set_condition(gamma, beta)
|
||||
|
||||
|
||||
@NORMALIZATION.register_module("SPADE")
|
||||
class SpatiallyAdaptiveDenormalization(AdaptiveDenormalization):
|
||||
def __init__(self, num_features: int, condition_in_channels, base_channels=128, base_norm_type="BN",
|
||||
activation_type="ReLU", padding_mode="zeros", gamma_bias=0.0):
|
||||
super().__init__(num_features, base_norm_type, gamma_bias=gamma_bias)
|
||||
self.base_conv_block = Conv2dBlock(condition_in_channels, base_channels, activation_type=activation_type,
|
||||
kernel_size=3, padding=1, padding_mode=padding_mode, norm_type="NONE")
|
||||
self.beta_conv = nn.Conv2d(base_channels, num_features, kernel_size=3, padding=1, padding_mode=padding_mode)
|
||||
self.gamma_conv = nn.Conv2d(base_channels, num_features, kernel_size=3, padding=1, padding_mode=padding_mode)
|
||||
|
||||
def set_condition_image(self, condition_image):
|
||||
feature = self.base_conv_block(condition_image)
|
||||
gamma = self.gamma_conv(feature)
|
||||
beta = self.beta_conv(feature)
|
||||
super().set_condition(gamma, beta)
|
||||
|
||||
|
||||
def _instance_layer_normalization(x, gamma, beta, rho, eps=1e-5):
|
||||
out = rho * F.instance_norm(x, eps=eps) + (1 - rho) * F.layer_norm(x, x.size()[1:], eps=eps)
|
||||
out = out * gamma.unsqueeze(2).unsqueeze(3) + beta.unsqueeze(2).unsqueeze(3)
|
||||
return out
|
||||
|
||||
|
||||
@NORMALIZATION.register_module("ILN")
|
||||
class ILN(nn.Module):
|
||||
def __init__(self, num_features, eps=1e-5):
|
||||
super(ILN, self).__init__()
|
||||
self.eps = eps
|
||||
self.rho = nn.Parameter(torch.Tensor(num_features))
|
||||
self.gamma = nn.Parameter(torch.Tensor(num_features))
|
||||
self.beta = nn.Parameter(torch.Tensor(num_features))
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
nn.init.zeros_(self.rho)
|
||||
nn.init.ones_(self.gamma)
|
||||
nn.init.zeros_(self.beta)
|
||||
|
||||
def forward(self, x):
|
||||
return _instance_layer_normalization(
|
||||
x, self.gamma.view(1, -1), self.beta.view(1, -1), self.rho.view(1, -1, 1, 1), self.eps)
|
||||
|
||||
|
||||
@NORMALIZATION.register_module("AdaILN")
|
||||
class AdaILN(nn.Module):
|
||||
def __init__(self, num_features, eps=1e-5, default_rho=0.9):
|
||||
super(AdaILN, self).__init__()
|
||||
self.eps = eps
|
||||
self.rho = nn.Parameter(torch.Tensor(num_features))
|
||||
self.rho.data.fill_(default_rho)
|
||||
|
||||
self.gamma = None
|
||||
self.beta = None
|
||||
self.have_set_condition = False
|
||||
|
||||
def set_condition(self, gamma, beta):
|
||||
self.gamma, self.beta = gamma, beta
|
||||
self.have_set_condition = True
|
||||
|
||||
def forward(self, x):
|
||||
assert self.have_set_condition
|
||||
out = _instance_layer_normalization(x, self.gamma, self.beta, self.rho.view(1, -1, 1, 1), self.eps)
|
||||
self.have_set_condition = False
|
||||
return out
|
||||
151
model/image_translation/CycleGAN.py
Normal file
151
model/image_translation/CycleGAN.py
Normal file
@@ -0,0 +1,151 @@
|
||||
import torch.nn as nn
|
||||
|
||||
from model import MODEL
|
||||
from model.base.module import Conv2dBlock, ResidualBlock
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, in_channels, base_channels, num_conv, num_res, max_down_sampling_multiple=2,
|
||||
padding_mode='reflect', activation_type="ReLU",
|
||||
down_conv_norm_type="IN", down_conv_kernel_size=3,
|
||||
res_norm_type="IN", pre_activation=False):
|
||||
super().__init__()
|
||||
|
||||
sequence = [Conv2dBlock(
|
||||
in_channels, base_channels, kernel_size=7, stride=1, padding=3, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=down_conv_norm_type
|
||||
)]
|
||||
multiple_now = 1
|
||||
for i in range(1, num_conv + 1):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** max_down_sampling_multiple)
|
||||
sequence.append(Conv2dBlock(
|
||||
multiple_prev * base_channels, multiple_now * base_channels,
|
||||
kernel_size=down_conv_kernel_size, stride=2, padding=1, padding_mode="zeros",
|
||||
activation_type=activation_type, norm_type=down_conv_norm_type
|
||||
))
|
||||
self.out_channels = multiple_now * base_channels
|
||||
sequence += [
|
||||
ResidualBlock(
|
||||
self.out_channels,
|
||||
padding_mode=padding_mode,
|
||||
activation_type=activation_type,
|
||||
norm_type=res_norm_type,
|
||||
pre_activation=pre_activation
|
||||
) for _ in range(num_res)
|
||||
]
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
|
||||
def forward(self, x):
|
||||
return self.sequence(x)
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, num_up_sampling, num_residual_blocks,
|
||||
activation_type="ReLU", padding_mode='reflect',
|
||||
up_conv_kernel_size=5, up_conv_norm_type="LN",
|
||||
res_norm_type="AdaIN", pre_activation=False, use_transpose_conv=False):
|
||||
super().__init__()
|
||||
self.residual_blocks = nn.ModuleList([
|
||||
ResidualBlock(
|
||||
in_channels,
|
||||
padding_mode=padding_mode,
|
||||
activation_type=activation_type,
|
||||
norm_type=res_norm_type,
|
||||
pre_activation=pre_activation
|
||||
) for _ in range(num_residual_blocks)
|
||||
])
|
||||
|
||||
sequence = list()
|
||||
channels = in_channels
|
||||
padding = (up_conv_kernel_size - 1) // 2
|
||||
for i in range(num_up_sampling):
|
||||
if use_transpose_conv:
|
||||
sequence.append(Conv2dBlock(
|
||||
channels, channels // 2, kernel_size=up_conv_kernel_size, stride=2,
|
||||
padding=padding, output_padding=padding,
|
||||
padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=up_conv_norm_type,
|
||||
use_transpose_conv=True
|
||||
))
|
||||
else:
|
||||
sequence.append(nn.Sequential(
|
||||
nn.Upsample(scale_factor=2),
|
||||
Conv2dBlock(channels, channels // 2, kernel_size=up_conv_kernel_size, stride=1,
|
||||
padding=padding, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=up_conv_norm_type),
|
||||
))
|
||||
channels = channels // 2
|
||||
sequence.append(Conv2dBlock(channels, out_channels, kernel_size=7, stride=1, padding=3,
|
||||
padding_mode=padding_mode, activation_type="Tanh", norm_type="NONE"))
|
||||
|
||||
self.up_sequence = nn.Sequential(*sequence)
|
||||
|
||||
def forward(self, x):
|
||||
for i, blk in enumerate(self.residual_blocks):
|
||||
x = blk(x)
|
||||
return self.up_sequence(x)
|
||||
|
||||
|
||||
@MODEL.register_module("CycleGAN-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, base_channels=64, num_blocks=9, activation_type="ReLU",
|
||||
padding_mode='reflect', norm_type="IN", pre_activation=False, use_transpose_conv=True):
|
||||
super().__init__()
|
||||
self.encoder = Encoder(in_channels, base_channels, num_conv=2, num_res=num_blocks,
|
||||
padding_mode=padding_mode, activation_type=activation_type,
|
||||
down_conv_norm_type=norm_type, res_norm_type=norm_type, pre_activation=pre_activation)
|
||||
self.decoder = Decoder(self.encoder.out_channels, out_channels, num_up_sampling=2, num_residual_blocks=0,
|
||||
padding_mode=padding_mode, activation_type=activation_type,
|
||||
up_conv_kernel_size=3, up_conv_norm_type=norm_type,
|
||||
pre_activation=pre_activation, use_transpose_conv=use_transpose_conv)
|
||||
|
||||
def forward(self, x):
|
||||
return self.decoder(self.encoder(x))
|
||||
|
||||
|
||||
@MODEL.register_module("PatchDiscriminator")
|
||||
class PatchDiscriminator(nn.Module):
|
||||
def __init__(self, in_channels, base_channels=64, num_conv=4, need_intermediate_feature=False,
|
||||
norm_type="IN", padding_mode='reflect', activation_type="LeakyReLU"):
|
||||
super().__init__()
|
||||
self.need_intermediate_feature = need_intermediate_feature
|
||||
kernel_size = 4
|
||||
padding = (kernel_size - 1) // 2
|
||||
sequence = [Conv2dBlock(
|
||||
in_channels, base_channels, kernel_size=kernel_size, stride=2, padding=padding, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=norm_type
|
||||
)]
|
||||
|
||||
multiple_now = 1
|
||||
for i in range(1, num_conv):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** 3)
|
||||
stride = 1 if i == num_conv - 1 else 2
|
||||
sequence.append(Conv2dBlock(
|
||||
multiple_prev * base_channels, multiple_now * base_channels,
|
||||
kernel_size=kernel_size, stride=stride, padding=padding, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=norm_type
|
||||
))
|
||||
sequence.append(nn.Conv2d(
|
||||
base_channels * multiple_now, 1, kernel_size, stride=1, padding=padding, padding_mode=padding_mode))
|
||||
if self.need_intermediate_feature:
|
||||
self.sequence = nn.ModuleList(sequence)
|
||||
else:
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
|
||||
def forward(self, x):
|
||||
if self.need_intermediate_feature:
|
||||
intermediate_feature = []
|
||||
for layer in self.sequence:
|
||||
x = layer(x)
|
||||
intermediate_feature.append(x)
|
||||
return tuple(intermediate_feature)
|
||||
else:
|
||||
return self.sequence(x)
|
||||
|
||||
if __name__ == '__main__':
|
||||
g = Generator(**dict(in_channels=3, out_channels=3))
|
||||
print(g)
|
||||
pd = PatchDiscriminator(**dict(in_channels=3, base_channels=64, num_conv=4))
|
||||
print(pd)
|
||||
183
model/image_translation/GauGAN.py
Normal file
183
model/image_translation/GauGAN.py
Normal file
@@ -0,0 +1,183 @@
|
||||
from collections import OrderedDict
|
||||
from functools import partial
|
||||
|
||||
import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from model.base.module import ResidualBlock, Conv2dBlock, LinearBlock
|
||||
from model import MODEL
|
||||
|
||||
class StyleEncoder(nn.Module):
|
||||
def __init__(self, in_channels, style_dim, num_conv, end_size=(4, 4), base_channels=64,
|
||||
norm_type="IN", padding_mode='reflect', activation_type="LeakyReLU"):
|
||||
super().__init__()
|
||||
sequence = [Conv2dBlock(
|
||||
in_channels, base_channels, kernel_size=3, stride=1, padding=1, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=norm_type
|
||||
)]
|
||||
multiple_now = 0
|
||||
max_multiple = 3
|
||||
for i in range(1, num_conv + 1):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** max_multiple)
|
||||
sequence.append(Conv2dBlock(
|
||||
multiple_prev * base_channels, multiple_now * base_channels,
|
||||
kernel_size=3, stride=2, padding=1, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=norm_type
|
||||
))
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
self.fc_avg = nn.Linear(base_channels * (2 ** max_multiple) * end_size[0] * end_size[1], style_dim)
|
||||
self.fc_var = nn.Linear(base_channels * (2 ** max_multiple) * end_size[0] * end_size[1], style_dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.sequence(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
return self.fc_avg(x), self.fc_var(x)
|
||||
|
||||
|
||||
class ImprovedSPADEGenerator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, output_size, have_style_input, style_dim, start_size=(4, 4),
|
||||
base_channels=64, padding_mode='reflect', activation_type="LeakyReLU", pre_activation=False):
|
||||
super().__init__()
|
||||
|
||||
assert output_size in (128, 256, 512, 1024)
|
||||
self.output_size = output_size
|
||||
|
||||
kernel_size = 3
|
||||
|
||||
if have_style_input:
|
||||
self.style_converter = nn.Sequential(
|
||||
LinearBlock(style_dim, 2 * style_dim, activation_type=activation_type, norm_type="NONE"),
|
||||
LinearBlock(2 * style_dim, 2 * style_dim, activation_type=activation_type, norm_type="NONE"),
|
||||
)
|
||||
|
||||
base_conv = partial(
|
||||
Conv2dBlock,
|
||||
pre_activation=pre_activation, activation_type=activation_type,
|
||||
norm_type="AdaIN" if have_style_input else "NONE",
|
||||
kernel_size=kernel_size, padding=(kernel_size - 1) // 2, padding_mode=padding_mode
|
||||
)
|
||||
|
||||
base_residual_block = partial(
|
||||
ResidualBlock,
|
||||
padding_mode=padding_mode,
|
||||
activation_type=activation_type,
|
||||
norm_type="SPADE",
|
||||
pre_activation=True,
|
||||
additional_norm_kwargs=dict(
|
||||
condition_in_channels=in_channels, base_channels=128, base_norm_type="BN",
|
||||
activation_type="ReLU", padding_mode="zeros", gamma_bias=1.0
|
||||
)
|
||||
)
|
||||
|
||||
sequence = OrderedDict()
|
||||
channels = (2 ** 4) * base_channels
|
||||
sequence["block_head"] = nn.Sequential(OrderedDict([
|
||||
("conv_input", base_conv(in_channels=in_channels, out_channels=channels)),
|
||||
("conv_style", base_conv(in_channels=channels, out_channels=channels)),
|
||||
("res_a", base_residual_block(in_channels=channels, out_channels=channels)),
|
||||
("res_b", base_residual_block(in_channels=channels, out_channels=channels)),
|
||||
("up", nn.Upsample(scale_factor=2, mode='nearest'))
|
||||
]))
|
||||
|
||||
for i in range(4, 9 - min(int(math.log(self.output_size, 2)), 8), -1):
|
||||
channels = (2 ** (i - 1)) * base_channels
|
||||
sequence[f"block_{2 * channels}"] = nn.Sequential(OrderedDict([
|
||||
("res_a", base_residual_block(in_channels=channels * 2, out_channels=channels)),
|
||||
("conv_style", base_conv(in_channels=channels, out_channels=channels)),
|
||||
("res_b", base_residual_block(in_channels=channels, out_channels=channels)),
|
||||
("up", nn.Upsample(scale_factor=2, mode='nearest'))
|
||||
]))
|
||||
self.sequence = nn.Sequential(sequence)
|
||||
# channels = 2*base_channels when output size is 256, 512, 1024
|
||||
# channels = 5*base_channels when output size is 128
|
||||
out_modules = OrderedDict()
|
||||
out_modules["out_1"] = nn.Sequential(
|
||||
Conv2dBlock(
|
||||
channels, out_channels, kernel_size=5, stride=1, padding=2,
|
||||
pre_activation=pre_activation,
|
||||
padding_mode=padding_mode, activation_type=activation_type, norm_type="NONE"
|
||||
),
|
||||
nn.Tanh()
|
||||
)
|
||||
for i in range(int(math.log(self.output_size, 2)) - 8):
|
||||
channels = channels // 2
|
||||
out_modules[f"block_{2 * channels}"] = nn.Sequential(OrderedDict([
|
||||
("res_a", base_residual_block(in_channels=2 * channels, out_channels=channels)),
|
||||
("res_b", base_residual_block(in_channels=channels, out_channels=channels)),
|
||||
("up", nn.Upsample(scale_factor=2, mode='nearest'))
|
||||
]))
|
||||
out_modules[f"out_{i + 2}"] = nn.Sequential(
|
||||
Conv2dBlock(
|
||||
channels, out_channels, kernel_size=5, stride=1, padding=2,
|
||||
pre_activation=pre_activation,
|
||||
padding_mode=padding_mode, activation_type=activation_type, norm_type="NONE"
|
||||
),
|
||||
nn.Tanh()
|
||||
)
|
||||
self.out_modules = nn.ModuleDict(out_modules)
|
||||
|
||||
def forward(self, seg, style=None):
|
||||
pass
|
||||
|
||||
@MODEL.register_module()
|
||||
class SPADEGenerator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, num_blocks, use_vae, num_z_dim, start_size=(4, 4), base_channels=64,
|
||||
padding_mode='reflect', activation_type="LeakyReLU"):
|
||||
super().__init__()
|
||||
self.sx, self.sy = start_size
|
||||
self.use_vae = use_vae
|
||||
self.num_z_dim = num_z_dim
|
||||
if use_vae:
|
||||
self.input_converter = nn.Linear(num_z_dim, 16 * base_channels * self.sx * self.sy)
|
||||
else:
|
||||
self.input_converter = nn.Conv2d(in_channels, 16 * base_channels, kernel_size=3, padding=1)
|
||||
|
||||
sequence = []
|
||||
|
||||
multiple_now = 16
|
||||
for i in range(num_blocks - 1, -1, -1):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** 4)
|
||||
if i != num_blocks - 1:
|
||||
sequence.append(nn.Upsample(scale_factor=2))
|
||||
sequence.append(ResidualBlock(
|
||||
base_channels * multiple_prev,
|
||||
out_channels=base_channels * multiple_now,
|
||||
padding_mode=padding_mode,
|
||||
activation_type=activation_type,
|
||||
norm_type="SPADE",
|
||||
pre_activation=True,
|
||||
additional_norm_kwargs=dict(
|
||||
condition_in_channels=in_channels, base_channels=128, base_norm_type="BN",
|
||||
activation_type="ReLU", padding_mode="zeros", gamma_bias=1.0
|
||||
)
|
||||
))
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
self.output_converter = Conv2dBlock(base_channels, out_channels, kernel_size=3, stride=1, padding=1,
|
||||
padding_mode=padding_mode, activation_type="Tanh", norm_type="NONE")
|
||||
|
||||
def forward(self, seg, z=None):
|
||||
if self.use_vae:
|
||||
if z is None:
|
||||
z = torch.randn(seg.size(0), self.num_z_dim, device=seg.device)
|
||||
x = self.input_converter(z).view(seg.size(0), -1, self.sx, self.sy)
|
||||
else:
|
||||
x = self.input_converter(F.interpolate(seg, size=(self.sx, self.sy)))
|
||||
for blk in self.sequence:
|
||||
if isinstance(blk, ResidualBlock):
|
||||
downsampling_seg = F.interpolate(seg, size=x.size()[2:], mode='nearest')
|
||||
blk.conv1.normalization.set_condition_image(downsampling_seg)
|
||||
blk.conv2.normalization.set_condition_image(downsampling_seg)
|
||||
if blk.learn_skip_connection:
|
||||
blk.res_conv.normalization.set_condition_image(downsampling_seg)
|
||||
x = blk(x)
|
||||
return self.output_converter(x)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
g = SPADEGenerator(3, 3, 7, False, 256)
|
||||
print(g)
|
||||
print(g(torch.randn(2, 3, 256, 256)).size())
|
||||
89
model/image_translation/MUNIT.py
Normal file
89
model/image_translation/MUNIT.py
Normal file
@@ -0,0 +1,89 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from model import MODEL
|
||||
from model.base.module import LinearBlock
|
||||
from model.image_translation.CycleGAN import Encoder, Decoder
|
||||
|
||||
|
||||
class StyleEncoder(nn.Module):
|
||||
def __init__(self, in_channels, out_dim, num_conv, base_channels=64,
|
||||
max_down_sampling_multiple=2, padding_mode='reflect', activation_type="ReLU", norm_type="NONE",
|
||||
pre_activation=False):
|
||||
super().__init__()
|
||||
self.down_encoder = Encoder(
|
||||
in_channels, base_channels, num_conv, num_res=0, max_down_sampling_multiple=max_down_sampling_multiple,
|
||||
padding_mode=padding_mode, activation_type=activation_type,
|
||||
down_conv_norm_type=norm_type, down_conv_kernel_size=4, pre_activation=pre_activation,
|
||||
)
|
||||
sequence = list()
|
||||
sequence.append(nn.AdaptiveAvgPool2d(1))
|
||||
# conv1x1 works as fc when tensor's size is (batch_size, channels, 1, 1), keep same with origin code
|
||||
sequence.append(nn.Conv2d(self.down_encoder.out_channels, out_dim, kernel_size=1, stride=1, padding=0))
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
|
||||
def forward(self, image):
|
||||
return self.sequence(image).view(image.size(0), -1)
|
||||
|
||||
|
||||
class MLPFusion(nn.Module):
|
||||
def __init__(self, in_features, out_features, base_features, n_blocks, activation_type="ReLU", norm_type="NONE"):
|
||||
super().__init__()
|
||||
|
||||
sequence = [LinearBlock(in_features, base_features, activation_type=activation_type, norm_type=norm_type)]
|
||||
sequence += [
|
||||
LinearBlock(base_features, base_features, activation_type=activation_type, norm_type=norm_type)
|
||||
for _ in range(n_blocks - 2)
|
||||
]
|
||||
sequence.append(LinearBlock(base_features, out_features, activation_type=activation_type, norm_type=norm_type))
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
|
||||
def forward(self, x):
|
||||
return self.sequence(x)
|
||||
|
||||
|
||||
@MODEL.register_module("MUNIT-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, base_channels=64, style_dim=8,
|
||||
num_mlp_base_feature=256, num_mlp_blocks=3,
|
||||
max_down_sampling_multiple=2, num_content_down_sampling=2, num_style_down_sampling=2,
|
||||
encoder_num_residual_blocks=4, decoder_num_residual_blocks=4,
|
||||
padding_mode='reflect', activation_type="ReLU", pre_activation=False):
|
||||
super().__init__()
|
||||
self.content_encoder = Encoder(
|
||||
in_channels, base_channels, num_content_down_sampling, encoder_num_residual_blocks,
|
||||
max_down_sampling_multiple=num_content_down_sampling,
|
||||
padding_mode=padding_mode, activation_type=activation_type,
|
||||
down_conv_norm_type="IN", down_conv_kernel_size=4,
|
||||
res_norm_type="IN", pre_activation=pre_activation
|
||||
)
|
||||
|
||||
self.style_encoder = StyleEncoder(in_channels, style_dim, num_style_down_sampling, base_channels,
|
||||
max_down_sampling_multiple, padding_mode, activation_type,
|
||||
norm_type="NONE", pre_activation=pre_activation)
|
||||
|
||||
content_channels = base_channels * (2 ** max_down_sampling_multiple)
|
||||
|
||||
self.fusion = MLPFusion(style_dim, decoder_num_residual_blocks * 2 * content_channels * 2,
|
||||
num_mlp_base_feature, num_mlp_blocks, activation_type,
|
||||
norm_type="NONE")
|
||||
|
||||
self.decoder = Decoder(in_channels, out_channels, max_down_sampling_multiple, decoder_num_residual_blocks,
|
||||
activation_type=activation_type, padding_mode=padding_mode,
|
||||
up_conv_kernel_size=5, up_conv_norm_type="LN",
|
||||
res_norm_type="AdaIN", pre_activation=pre_activation)
|
||||
|
||||
def encode(self, x):
|
||||
return self.content_encoder(x), self.style_encoder(x)
|
||||
|
||||
def decode(self, content, style):
|
||||
as_param_style = torch.chunk(self.fusion(style), 2 * len(self.decoder.residual_blocks), dim=1)
|
||||
# set style for decoder
|
||||
for i, blk in enumerate(self.decoder.residual_blocks):
|
||||
blk.conv1.normalization.set_style(as_param_style[2 * i])
|
||||
blk.conv2.normalization.set_style(as_param_style[2 * i + 1])
|
||||
self.decoder(content)
|
||||
|
||||
def forward(self, x):
|
||||
content, style = self.encode(x)
|
||||
return self.decode(content, style)
|
||||
98
model/image_translation/TSIT.py
Normal file
98
model/image_translation/TSIT.py
Normal file
@@ -0,0 +1,98 @@
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
|
||||
from model import MODEL
|
||||
from model.base.module import ResidualBlock, Conv2dBlock
|
||||
|
||||
|
||||
class Interpolation(nn.Module):
|
||||
def __init__(self, scale_factor=None, mode='nearest', align_corners=None):
|
||||
super(Interpolation, self).__init__()
|
||||
self.scale_factor = scale_factor
|
||||
self.mode = mode
|
||||
self.align_corners = align_corners
|
||||
|
||||
def forward(self, x):
|
||||
return F.interpolate(x, scale_factor=self.scale_factor, mode=self.mode, align_corners=self.align_corners,
|
||||
recompute_scale_factor=False)
|
||||
|
||||
def __repr__(self):
|
||||
return f"Interpolation(scale_factor={self.scale_factor}, mode={self.mode}, align_corners={self.align_corners})"
|
||||
|
||||
|
||||
@MODEL.register_module("TSIT-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, base_channels=64, num_blocks=7,
|
||||
padding_mode='reflect', activation_type="LeakyReLU"):
|
||||
super().__init__()
|
||||
self.input_layer = Conv2dBlock(
|
||||
in_channels, base_channels, kernel_size=7, stride=1, padding=3, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type="IN",
|
||||
)
|
||||
multiple_now = 1
|
||||
stream_sequence = []
|
||||
for i in range(1, num_blocks + 1):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** 4)
|
||||
stream_sequence.append(nn.Sequential(
|
||||
Interpolation(scale_factor=0.5, mode="nearest"),
|
||||
ResidualBlock(
|
||||
multiple_prev * base_channels, out_channels=multiple_now * base_channels,
|
||||
padding_mode=padding_mode, activation_type=activation_type, norm_type="IN")
|
||||
))
|
||||
self.down_sequence = nn.ModuleList(stream_sequence)
|
||||
|
||||
|
||||
sequence = []
|
||||
multiple_now = 16
|
||||
for i in range(num_blocks - 1, -1, -1):
|
||||
multiple_prev = multiple_now
|
||||
multiple_now = min(2 ** i, 2 ** 4)
|
||||
sequence.append(nn.Sequential(
|
||||
ResidualBlock(
|
||||
base_channels * multiple_prev,
|
||||
out_channels=base_channels * multiple_now,
|
||||
padding_mode=padding_mode,
|
||||
activation_type=activation_type,
|
||||
norm_type="FADE",
|
||||
pre_activation=True,
|
||||
additional_norm_kwargs=dict(
|
||||
condition_in_channels=base_channels * multiple_prev, base_norm_type="BN",
|
||||
padding_mode="zeros", gamma_bias=0.0
|
||||
)
|
||||
),
|
||||
Interpolation(scale_factor=2, mode="nearest")
|
||||
))
|
||||
self.up_sequence = nn.Sequential(*sequence)
|
||||
|
||||
self.output_layer = Conv2dBlock(
|
||||
base_channels, out_channels, kernel_size=3, stride=1, padding=1,
|
||||
padding_mode=padding_mode, activation_type="Tanh", norm_type="NONE"
|
||||
)
|
||||
|
||||
def forward(self, x, z=None):
|
||||
c = self.input_layer(x)
|
||||
contents = []
|
||||
for blk in self.down_sequence:
|
||||
c = blk(c)
|
||||
contents.append(c)
|
||||
if z is None:
|
||||
# for image 256x256, z size: [batch_size, 1024, 2, 2]
|
||||
z = torch.randn(size=contents[-1].size(), device=contents[-1].device)
|
||||
|
||||
for blk in self.up_sequence:
|
||||
res = blk[0]
|
||||
c = contents.pop()
|
||||
res.conv1.normalization.set_feature(c)
|
||||
res.conv2.normalization.set_feature(c)
|
||||
if res.learn_skip_connection:
|
||||
res.res_conv.normalization.set_feature(c)
|
||||
return self.output_layer(self.up_sequence(z))
|
||||
|
||||
if __name__ == '__main__':
|
||||
g = Generator(3, 3).cuda()
|
||||
img = torch.randn(2, 3, 256, 256).cuda()
|
||||
print(g(img).size())
|
||||
|
||||
|
||||
125
model/image_translation/UGATIT.py
Normal file
125
model/image_translation/UGATIT.py
Normal file
@@ -0,0 +1,125 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from model import MODEL
|
||||
from model.base.module import Conv2dBlock, LinearBlock
|
||||
from model.image_translation.CycleGAN import Encoder, Decoder
|
||||
|
||||
|
||||
class CAMClassifier(nn.Module):
|
||||
def __init__(self, in_channels, activation_type="ReLU"):
|
||||
super(CAMClassifier, self).__init__()
|
||||
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.avg_fc = nn.Linear(in_channels, 1, bias=False)
|
||||
self.max_pool = nn.AdaptiveMaxPool2d(1)
|
||||
self.max_fc = nn.Linear(in_channels, 1, bias=False)
|
||||
self.fusion_conv = Conv2dBlock(in_channels * 2, in_channels, kernel_size=1, stride=1, bias=True,
|
||||
activation_type=activation_type, norm_type="NONE")
|
||||
|
||||
def forward(self, x):
|
||||
avg_logit = self.avg_fc(self.avg_pool(x).view(x.size(0), -1))
|
||||
max_logit = self.max_fc(self.max_pool(x).view(x.size(0), -1))
|
||||
|
||||
return self.fusion_conv(torch.cat(
|
||||
[x * self.avg_fc.weight.unsqueeze(2).unsqueeze(3), x * self.max_fc.weight.unsqueeze(2).unsqueeze(3)],
|
||||
dim=1
|
||||
)), torch.cat([avg_logit, max_logit], 1)
|
||||
|
||||
|
||||
@MODEL.register_module("UGATIT-Generator")
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, base_channels=64, num_blocks=6, img_size=256, light=False,
|
||||
activation_type="ReLU", norm_type="IN", padding_mode='reflect', pre_activation=False):
|
||||
super(Generator, self).__init__()
|
||||
|
||||
self.light = light
|
||||
|
||||
n_down_sampling = 2
|
||||
self.encoder = Encoder(in_channels, base_channels, n_down_sampling, num_blocks,
|
||||
padding_mode=padding_mode, activation_type=activation_type,
|
||||
down_conv_norm_type=norm_type, down_conv_kernel_size=3, res_norm_type=norm_type,
|
||||
pre_activation=pre_activation)
|
||||
mult = 2 ** n_down_sampling
|
||||
self.cam = CAMClassifier(base_channels * mult, activation_type)
|
||||
|
||||
# Gamma, Beta block
|
||||
if self.light:
|
||||
self.fc = nn.Sequential(
|
||||
LinearBlock(base_channels * mult, base_channels * mult, False, "ReLU", "NONE"),
|
||||
LinearBlock(base_channels * mult, base_channels * mult, False, "ReLU", "NONE")
|
||||
)
|
||||
else:
|
||||
self.fc = nn.Sequential(
|
||||
LinearBlock(img_size // mult * img_size // mult * base_channels * mult, base_channels * mult, False,
|
||||
"ReLU", "NONE"),
|
||||
LinearBlock(base_channels * mult, base_channels * mult, False, "ReLU", "NONE")
|
||||
)
|
||||
|
||||
self.gamma = nn.Linear(base_channels * mult, base_channels * mult, bias=False)
|
||||
self.beta = nn.Linear(base_channels * mult, base_channels * mult, bias=False)
|
||||
|
||||
self.decoder = Decoder(
|
||||
base_channels * mult, out_channels, n_down_sampling, num_blocks,
|
||||
activation_type=activation_type, padding_mode=padding_mode,
|
||||
up_conv_kernel_size=3, up_conv_norm_type="ILN",
|
||||
res_norm_type="AdaILN", pre_activation=pre_activation
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.encoder(x)
|
||||
|
||||
x, cam_logit = self.cam(x)
|
||||
|
||||
heatmap = torch.sum(x, dim=1, keepdim=True)
|
||||
|
||||
if self.light:
|
||||
x_ = torch.nn.functional.adaptive_avg_pool2d(x, (1, 1))
|
||||
x_ = self.fc(x_.view(x_.shape[0], -1))
|
||||
else:
|
||||
x_ = self.fc(x.view(x.shape[0], -1))
|
||||
gamma, beta = self.gamma(x_), self.beta(x_)
|
||||
|
||||
for blk in self.decoder.residual_blocks:
|
||||
blk.conv1.normalization.set_condition(gamma, beta)
|
||||
blk.conv2.normalization.set_condition(gamma, beta)
|
||||
return self.decoder(x), cam_logit, heatmap
|
||||
|
||||
|
||||
@MODEL.register_module("UGATIT-Discriminator")
|
||||
class Discriminator(nn.Module):
|
||||
def __init__(self, in_channels, base_channels=64, num_blocks=5,
|
||||
activation_type="LeakyReLU", norm_type="NONE", padding_mode='reflect'):
|
||||
super().__init__()
|
||||
|
||||
sequence = [Conv2dBlock(
|
||||
in_channels, base_channels, kernel_size=4, stride=2, padding=1, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=norm_type
|
||||
)]
|
||||
|
||||
sequence += [Conv2dBlock(
|
||||
base_channels * (2 ** i), base_channels * (2 ** i) * 2,
|
||||
kernel_size=4, stride=2, padding=1, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=norm_type) for i in range(num_blocks - 3)]
|
||||
|
||||
sequence.append(
|
||||
Conv2dBlock(base_channels * (2 ** (num_blocks - 3)), base_channels * (2 ** (num_blocks - 2)),
|
||||
kernel_size=4, stride=1, padding=1, padding_mode=padding_mode,
|
||||
activation_type=activation_type, norm_type=norm_type)
|
||||
)
|
||||
self.sequence = nn.Sequential(*sequence)
|
||||
|
||||
mult = 2 ** (num_blocks - 2)
|
||||
self.cam = CAMClassifier(base_channels * mult, activation_type)
|
||||
self.conv = nn.Conv2d(base_channels * mult, 1, kernel_size=4, stride=1, padding=1, bias=False,
|
||||
padding_mode="reflect")
|
||||
|
||||
def forward(self, x, return_heatmap=False):
|
||||
x = self.sequence(x)
|
||||
|
||||
x, cam_logit = self.cam(x)
|
||||
|
||||
if return_heatmap:
|
||||
heatmap = torch.sum(x, dim=1, keepdim=True)
|
||||
return self.conv(x), cam_logit, heatmap
|
||||
else:
|
||||
return self.conv(x), cam_logit
|
||||
0
model/image_translation/__init__.py
Normal file
0
model/image_translation/__init__.py
Normal file
@@ -6,16 +6,20 @@ from model import MODEL
|
||||
|
||||
@MODEL.register_module()
|
||||
class MultiScaleDiscriminator(nn.Module):
|
||||
def __init__(self, num_scale, discriminator_cfg):
|
||||
def __init__(self, num_scale, discriminator_cfg, down_sample_method="avg"):
|
||||
super().__init__()
|
||||
assert down_sample_method in ["avg", "bilinear"]
|
||||
self.down_sample_method = down_sample_method
|
||||
|
||||
self.discriminator_list = nn.ModuleList([
|
||||
MODEL.build_with(discriminator_cfg) for _ in range(num_scale)
|
||||
])
|
||||
|
||||
@staticmethod
|
||||
def down_sample(x):
|
||||
return F.avg_pool2d(x, kernel_size=3, stride=2, padding=[1, 1], count_include_pad=False)
|
||||
def down_sample(self, x):
|
||||
if self.down_sample_method == "avg":
|
||||
return F.avg_pool2d(x, kernel_size=3, stride=2, padding=[1, 1], count_include_pad=False)
|
||||
if self.down_sample_method == "bilinear":
|
||||
return F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=True)
|
||||
|
||||
def forward(self, x):
|
||||
results = []
|
||||
@@ -1,76 +0,0 @@
|
||||
import functools
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def select_norm_layer(norm_type):
|
||||
if norm_type == "BN":
|
||||
return functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
|
||||
elif norm_type == "IN":
|
||||
return functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
|
||||
elif norm_type == "LN":
|
||||
return functools.partial(LayerNorm2d, affine=True)
|
||||
elif norm_type == "NONE":
|
||||
return lambda num_features: nn.Identity()
|
||||
elif norm_type == "AdaIN":
|
||||
return functools.partial(AdaptiveInstanceNorm2d, affine=False, track_running_stats=False)
|
||||
else:
|
||||
raise NotImplemented(f'normalization layer {norm_type} is not found')
|
||||
|
||||
|
||||
class LayerNorm2d(nn.Module):
|
||||
def __init__(self, num_features, eps: float = 1e-5, affine: bool = True):
|
||||
super().__init__()
|
||||
self.num_features = num_features
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
if self.affine:
|
||||
self.channel_gamma = nn.Parameter(torch.Tensor(1, num_features, 1, 1))
|
||||
self.channel_beta = nn.Parameter(torch.Tensor(1, num_features, 1, 1))
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
if self.affine:
|
||||
nn.init.uniform_(self.channel_gamma)
|
||||
nn.init.zeros_(self.channel_beta)
|
||||
|
||||
def forward(self, x):
|
||||
ln_mean, ln_var = torch.mean(x, dim=[1, 2, 3], keepdim=True), torch.var(x, dim=[1, 2, 3], keepdim=True)
|
||||
x = (x - ln_mean) / torch.sqrt(ln_var + self.eps)
|
||||
if self.affine:
|
||||
return self.channel_gamma * x + self.channel_beta
|
||||
return x
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.__class__.__name__}(num_features={self.num_features}, affine={self.affine})"
|
||||
|
||||
|
||||
class AdaptiveInstanceNorm2d(nn.Module):
|
||||
def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1,
|
||||
affine: bool = False, track_running_stats: bool = False):
|
||||
super().__init__()
|
||||
self.num_features = num_features
|
||||
self.affine = affine
|
||||
self.track_running_stats = track_running_stats
|
||||
self.norm = nn.InstanceNorm2d(num_features, eps, momentum, affine, track_running_stats)
|
||||
|
||||
self.gamma = None
|
||||
self.beta = None
|
||||
self.have_set_style = False
|
||||
|
||||
def set_style(self, style):
|
||||
style = style.view(*style.size(), 1, 1)
|
||||
self.gamma, self.beta = style.chunk(2, 1)
|
||||
self.have_set_style = True
|
||||
|
||||
def forward(self, x):
|
||||
assert self.have_set_style
|
||||
x = self.norm(x)
|
||||
x = self.gamma * x + self.beta
|
||||
self.have_set_style = False
|
||||
return x
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.__class__.__name__}(num_features={self.num_features}, " \
|
||||
f"affine={self.affine}, track_running_stats={self.track_running_stats})"
|
||||
@@ -1,3 +1,4 @@
|
||||
from util.registry import Registry
|
||||
|
||||
MODEL = Registry("model")
|
||||
NORMALIZATION = Registry("normalization")
|
||||
@@ -65,7 +65,8 @@ def generation_init_weights(module, init_type='normal', init_gain=0.02):
|
||||
elif classname.find('BatchNorm2d') != -1:
|
||||
# BatchNorm Layer's weight is not a matrix;
|
||||
# only normal distribution applies.
|
||||
normal_init(m, 1.0, init_gain)
|
||||
if m.weight is not None:
|
||||
normal_init(m, 1.0, init_gain)
|
||||
|
||||
assert isinstance(module, nn.Module)
|
||||
module.apply(init_func)
|
||||
|
||||
14
tool/inspect_model.py
Normal file
14
tool/inspect_model.py
Normal file
@@ -0,0 +1,14 @@
|
||||
import sys
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from engine.util.build import build_model
|
||||
|
||||
config = OmegaConf.load(sys.argv[1])
|
||||
|
||||
|
||||
generator = build_model(config.model.generator)
|
||||
|
||||
ckp = torch.load(sys.argv[2], map_location="cpu")
|
||||
|
||||
generator.module.load_state_dict(ckp["generator_main"])
|
||||
13
tool/process/permutation_face.py
Normal file
13
tool/process/permutation_face.py
Normal file
@@ -0,0 +1,13 @@
|
||||
from pathlib import Path
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from itertools import permutations
|
||||
|
||||
pids = defaultdict(list)
|
||||
for p in Path(sys.argv[1]).glob("*.jpg"):
|
||||
pids[p.stem[:7]].append(p.stem)
|
||||
|
||||
data = []
|
||||
for p in pids:
|
||||
data.extend(list(permutations(pids[p], 2)))
|
||||
|
||||
19
util/misc.py
19
util/misc.py
@@ -1,7 +1,26 @@
|
||||
import importlib
|
||||
import logging
|
||||
import pkgutil
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
def import_submodules(package, recursive=True):
|
||||
""" Import all submodules of a module, recursively, including subpackages
|
||||
|
||||
:param package: package (name or actual module)
|
||||
:type package: str | module
|
||||
:rtype: dict[str, types.ModuleType]
|
||||
"""
|
||||
if isinstance(package, str):
|
||||
package = importlib.import_module(package)
|
||||
results = {}
|
||||
for loader, name, is_pkg in pkgutil.walk_packages(package.__path__):
|
||||
full_name = package.__name__ + '.' + name
|
||||
results[name] = importlib.import_module(full_name)
|
||||
if recursive and is_pkg:
|
||||
results.update(import_submodules(full_name))
|
||||
return results
|
||||
|
||||
|
||||
def setup_logger(
|
||||
name: Optional[str] = None,
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
import inspect
|
||||
from omegaconf.dictconfig import DictConfig
|
||||
from omegaconf import OmegaConf
|
||||
from types import ModuleType
|
||||
import warnings
|
||||
from types import ModuleType
|
||||
|
||||
from omegaconf import OmegaConf
|
||||
from omegaconf.dictconfig import DictConfig
|
||||
|
||||
|
||||
class _Registry:
|
||||
def __init__(self, name):
|
||||
@@ -51,11 +53,9 @@ class _Registry:
|
||||
else:
|
||||
raise TypeError(f'cfg must be a dict or a str, but got {type(cfg)}')
|
||||
|
||||
for k in args:
|
||||
assert isinstance(k, str)
|
||||
if k.startswith("_"):
|
||||
warnings.warn(f"got param start with `_`: {k}, will remove it")
|
||||
args.pop(k)
|
||||
for invalid_key in [k for k in args.keys() if k.startswith("_")]:
|
||||
warnings.warn(f"got param start with `_`: {invalid_key}, will remove it")
|
||||
args.pop(invalid_key)
|
||||
|
||||
if not (isinstance(default_args, dict) or default_args is None):
|
||||
raise TypeError('default_args must be a dict or None, '
|
||||
@@ -136,8 +136,11 @@ class Registry(_Registry):
|
||||
if module_name is None:
|
||||
module_name = module_class.__name__
|
||||
if not force and module_name in self._module_dict:
|
||||
raise KeyError(f'{module_name} is already registered '
|
||||
f'in {self.name}')
|
||||
if self._module_dict[module_name] == module_class:
|
||||
warnings.warn(f'{module_name} is already registered in {self.name}, but is the same class')
|
||||
return
|
||||
raise KeyError(f'{module_name}:{self._module_dict[module_name]} is already registered in {self.name}'
|
||||
f'so {module_class} can not be registered')
|
||||
self._module_dict[module_name] = module_class
|
||||
|
||||
def register_module(self, name=None, force=False, module=None):
|
||||
|
||||
Reference in New Issue
Block a user