File size: 1,809 Bytes
4eef8da | 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 |
import torch
import torch.nn as nn
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(
in_channels,
out_channels,
kernel_size=3,
padding=1
),
nn.ReLU(inplace=True),
nn.Conv2d(
out_channels,
out_channels,
kernel_size=3,
padding=1
),
nn.ReLU(inplace=True)
)
def forward(self, x):
return self.block(x)
class Encoder(nn.Module):
def __init__(self):
super().__init__()
self.stage_1 = ConvBlock(
in_channels=3,
out_channels=9
)
self.stage_2 = ConvBlock(
in_channels=9,
out_channels=27
)
self.stage_3 = nn.Conv2d(
in_channels=27,
out_channels=3,
kernel_size=3,
padding=1
)
def forward(self, x):
x = self.stage_1(x)
x = self.stage_2(x)
x = self.stage_3(x)
return x
class Decoder(nn.Module):
def __init__(self):
super().__init__()
self.upsample_block = nn.Upsample(
size=(640, 1024),
mode="bilinear",
align_corners=False
)
def forward(self, x):
return self.upsample_block(x)
class ImageRegressionNet(nn.Module):
def __init__(self):
super().__init__()
self.encoder = Encoder()
self.decoder = Decoder()
self.encoder_decoder = nn.Sequential(
self.encoder,
self.decoder
)
def forward(self, x):
return self.encoder_decoder(x)
|