File size: 1,282 Bytes
2a13495 |
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 |
import torch.nn as nn
from .modules import Activation
class SegmentationHead(nn.Sequential):
def __init__(
self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1
):
conv2d = nn.Conv2d(
in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2
)
upsampling = (
nn.UpsamplingBilinear2d(scale_factor=upsampling)
if upsampling > 1
else nn.Identity()
)
activation = Activation(activation)
super().__init__(conv2d, upsampling, activation)
class ClassificationHead(nn.Sequential):
def __init__(
self, in_channels, classes, pooling="avg", dropout=0.2, activation=None
):
if pooling not in ("max", "avg"):
raise ValueError(
"Pooling should be one of ('max', 'avg'), got {}.".format(pooling)
)
pool = nn.AdaptiveAvgPool2d(1) if pooling == "avg" else nn.AdaptiveMaxPool2d(1)
flatten = nn.Flatten()
dropout = nn.Dropout(p=dropout, inplace=True) if dropout else nn.Identity()
linear = nn.Linear(in_channels, classes, bias=True)
activation = Activation(activation)
super().__init__(pool, flatten, dropout, linear, activation)
|