mirror of
https://github.com/jantic/DeOldify.git
synced 2026-08-30 18:02:24 +08:00
Alexandre Vicenzi
547fb7e56f
It will be better to do:
from deoldify import visualize
Than:
from fasterai import visualize
The PyPI package will be called DeOldify, in this case
it makes more sense to have an import name that matches
the package name.
Also, fasterai resembles fastai library, and DeOldify
is not a library to use with fastai, it's built on
top of it.
25 lines
1.4 KiB
Python
25 lines
1.4 KiB
Python
from fastai.layers import *
|
|
from fastai.torch_core import *
|
|
from torch.nn.parameter import Parameter
|
|
from torch.autograd import Variable
|
|
|
|
|
|
#The code below is meant to be merged into fastaiv1 ideally
|
|
|
|
def custom_conv_layer(ni:int, nf:int, ks:int=3, stride:int=1, padding:int=None, bias:bool=None, is_1d:bool=False,
|
|
norm_type:Optional[NormType]=NormType.Batch, use_activ:bool=True, leaky:float=None,
|
|
transpose:bool=False, init:Callable=nn.init.kaiming_normal_, self_attention:bool=False,
|
|
extra_bn:bool=False):
|
|
"Create a sequence of convolutional (`ni` to `nf`), ReLU (if `use_activ`) and batchnorm (if `bn`) layers."
|
|
if padding is None: padding = (ks-1)//2 if not transpose else 0
|
|
bn = norm_type in (NormType.Batch, NormType.BatchZero) or extra_bn==True
|
|
if bias is None: bias = not bn
|
|
conv_func = nn.ConvTranspose2d if transpose else nn.Conv1d if is_1d else nn.Conv2d
|
|
conv = init_default(conv_func(ni, nf, kernel_size=ks, bias=bias, stride=stride, padding=padding), init)
|
|
if norm_type==NormType.Weight: conv = weight_norm(conv)
|
|
elif norm_type==NormType.Spectral: conv = spectral_norm(conv)
|
|
layers = [conv]
|
|
if use_activ: layers.append(relu(True, leaky=leaky))
|
|
if bn: layers.append((nn.BatchNorm1d if is_1d else nn.BatchNorm2d)(nf))
|
|
if self_attention: layers.append(SelfAttention(nf))
|
|
return nn.Sequential(*layers) |