forked from fastai/fastai
/
layers.py
126 lines (105 loc) · 5.73 KB
/
layers.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
"`fastai.layers` provides essential functions to building and modifying `model` architectures"
from .torch_core import *
__all__ = ['AdaptiveConcatPool2d', 'MSELossFlat', 'CrossEntropyFlat', 'Debugger', 'Flatten', 'Lambda', 'PoolFlatten', 'ResizeBatch',
'StdUpsample', 'bn_drop_lin', 'conv2d', 'conv2d_relu', 'conv2d_trans', 'conv_layer', 'get_embedding', 'simple_cnn',
'std_upsample_head', 'trunc_normal_']
class Lambda(nn.Module):
"An easy way to create a pytorch layer for a simple `func`."
def __init__(self, func:LambdaFunc):
"create a layer that simply calls `func` with `x`"
super().__init__()
self.func=func
def forward(self, x): return self.func(x)
def ResizeBatch(*size:int) -> Tensor:
"Layer that resizes x to `size`, good for connecting mismatched layers."
return Lambda(lambda x: x.view((-1,)+size))
def Flatten()->Tensor:
"Flattens `x` to a single dimension, often used at the end of a model."
return Lambda(lambda x: x.view((x.size(0), -1)))
def PoolFlatten()->nn.Sequential:
"Apply `nn.AdaptiveAvgPool2d` to `x` and then flatten the result."
return nn.Sequential(nn.AdaptiveAvgPool2d(1), Flatten())
def bn_drop_lin(n_in:int, n_out:int, bn:bool=True, p:float=0., actn:Optional[nn.Module]=None):
"`n_in`->bn->dropout->linear(`n_in`,`n_out`)->`actn`"
layers = [nn.BatchNorm1d(n_in)] if bn else []
if p != 0: layers.append(nn.Dropout(p))
layers.append(nn.Linear(n_in, n_out))
if actn is not None: layers.append(actn)
return layers
def conv2d(ni:int, nf:int, ks:int=3, stride:int=1, padding:int=None, bias=False) -> nn.Conv2d:
"Create `nn.Conv2d` layer: `ni` inputs, `nf` outputs, `ks` kernel size. `padding` defaults to `k//2`."
if padding is None: padding = ks//2
return nn.Conv2d(ni, nf, kernel_size=ks, stride=stride, padding=padding, bias=bias)
def conv_layer(ni:int, nf:int, ks:int=3, stride:int=1)->nn.Sequential:
"Create Conv2d->BatchNorm2d->LeakyReLu layer: `ni` input, `nf` out filters, `ks` kernel, `stride`:stride."
return nn.Sequential(
nn.Conv2d(ni, nf, kernel_size=ks, bias=False, stride=stride, padding=ks//2),
nn.BatchNorm2d(nf),
nn.LeakyReLU(negative_slope=0.1, inplace=True))
def conv2d_relu(ni:int, nf:int, ks:int=3, stride:int=1, padding:int=None, bn:bool=False,
bias:bool=False) -> nn.Sequential:
"""Create a `conv2d` layer with `nn.ReLU` activation and optional(`bn`) `nn.BatchNorm2d`: `ni` input, `nf` out
filters, `ks` kernel, `stride`:stride, `padding`:padding, `bn`: batch normalization."""
layers = [conv2d(ni, nf, ks=ks, stride=stride, padding=padding, bias=bias), nn.ReLU(inplace=True)]
if bn: layers.append(nn.BatchNorm2d(nf))
return nn.Sequential(*layers)
def conv2d_trans(ni:int, nf:int, ks:int=2, stride:int=2, padding:int=0, bias=False) -> nn.ConvTranspose2d:
"Create `nn.ConvTranspose2d` layer: `ni` inputs, `nf` outputs, `ks` kernel size, `stride`: stride. `padding` defaults to 0."
return nn.ConvTranspose2d(ni, nf, kernel_size=ks, stride=stride, padding=padding, bias=bias)
class AdaptiveConcatPool2d(nn.Module):
"Layer that concats `AdaptiveAvgPool2d` and `AdaptiveMaxPool2d`."
def __init__(self, sz:Optional[int]=None):
"Output will be 2*sz or 2 if sz is None"
super().__init__()
sz = sz or 1
self.ap,self.mp = nn.AdaptiveAvgPool2d(sz), nn.AdaptiveMaxPool2d(sz)
def forward(self, x): return torch.cat([self.mp(x), self.ap(x)], 1)
class Debugger(nn.Module):
"A module to debug inside a model."
def forward(self,x:Tensor) -> Tensor:
set_trace()
return x
class StdUpsample(nn.Module):
"Increases the dimensionality of our data by applying a transposed convolution layer."
def __init__(self, n_in:int, n_out:int):
super().__init__()
self.conv = conv2d_trans(n_in, n_out)
self.bn = nn.BatchNorm2d(n_out)
def forward(self, x:Tensor) -> Tensor:
return self.bn(F.relu(self.conv(x)))
def std_upsample_head(c, *nfs:Collection[int]) -> nn.Module:
"Create a sequence of upsample layers."
return nn.Sequential(
nn.ReLU(),
*(StdUpsample(nfs[i],nfs[i+1]) for i in range(4)),
conv2d_trans(nfs[-1], c)
)
class CrossEntropyFlat(nn.CrossEntropyLoss):
"Same as `nn.CrossEntropyLoss`, but flattens input and target."
def forward(self, input:Tensor, target:Tensor) -> Rank0Tensor:
n,c,*_ = input.shape
return super().forward(input.view(n, c, -1), target.view(n, -1))
class MSELossFlat(nn.MSELoss):
"Same as `nn.MSELoss`, but flattens input and target."
def forward(self, input:Tensor, target:Tensor) -> Rank0Tensor:
return super().forward(input.view(-1), target.view(-1))
def simple_cnn(actns:Collection[int], kernel_szs:Collection[int]=None,
strides:Collection[int]=None, bn=False) -> nn.Sequential:
"CNN with `conv2d_relu` layers defined by `actns`, `kernel_szs` and `strides`, plus batchnorm if `bn`."
nl = len(actns)-1
kernel_szs = ifnone(kernel_szs, [3]*nl)
strides = ifnone(strides , [2]*nl)
layers = [conv2d_relu(actns[i], actns[i+1], kernel_szs[i], stride=strides[i],
bn=(bn and i<(len(strides)-1))) for i in range_of(strides)]
layers.append(PoolFlatten())
return nn.Sequential(*layers)
def trunc_normal_(x:Tensor, mean:float=0., std:float=1.) -> Tensor:
"Truncated normal initialization."
# From https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/12
return x.normal_().fmod_(2).mul_(std).add_(mean)
def get_embedding(ni:int,nf:int) -> nn.Module:
"Create an embedding layer."
emb = nn.Embedding(ni, nf)
# See https://arxiv.org/abs/1711.09160
with torch.no_grad(): trunc_normal_(emb.weight, std=0.01)
return emb