mirror of
https://github.com/jantic/DeOldify.git
synced 2026-09-17 17:19:24 +08:00
Alexandre Vicenzi
3fb4956fdd
Black was used to format the code as well some pylint issues where fixed.
45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
from fastai.core import *
|
|
from fastai.torch_core import *
|
|
from fastai.vision import *
|
|
from fastai.vision.gan import AdaptiveLoss, accuracy_thresh_expand
|
|
|
|
_conv_args = dict(leaky=0.2, norm_type=NormType.Spectral)
|
|
|
|
|
|
def _conv(ni: int, nf: int, ks: int = 3, stride: int = 1, **kwargs):
|
|
return conv_layer(ni, nf, ks=ks, stride=stride, **_conv_args, **kwargs)
|
|
|
|
|
|
def custom_gan_critic(
|
|
n_channels: int = 3, nf: int = 256, n_blocks: int = 3, p: int = 0.15
|
|
):
|
|
"Critic to train a `GAN`."
|
|
layers = [_conv(n_channels, nf, ks=4, stride=2), nn.Dropout2d(p / 2)]
|
|
for i in range(n_blocks):
|
|
layers += [
|
|
_conv(nf, nf, ks=3, stride=1),
|
|
nn.Dropout2d(p),
|
|
_conv(nf, nf * 2, ks=4, stride=2, self_attention=(i == 0)),
|
|
]
|
|
nf *= 2
|
|
layers += [
|
|
_conv(nf, nf, ks=3, stride=1),
|
|
_conv(nf, 1, ks=4, bias=False, padding=0, use_activ=False),
|
|
Flatten(),
|
|
]
|
|
return nn.Sequential(*layers)
|
|
|
|
|
|
def colorize_crit_learner(
|
|
data: ImageDataBunch,
|
|
loss_critic=AdaptiveLoss(nn.BCEWithLogitsLoss()),
|
|
nf: int = 256,
|
|
) -> Learner:
|
|
return Learner(
|
|
data,
|
|
custom_gan_critic(nf=nf),
|
|
metrics=accuracy_thresh_expand,
|
|
loss_func=loss_critic,
|
|
wd=1e-3,
|
|
)
|