Spaces:
Sleeping
Sleeping
File size: 2,729 Bytes
302b2b5 237774d 302b2b5 237774d 302b2b5 237774d d57909e 237774d 302b2b5 d57909e 237774d 302b2b5 237774d d57909e 237774d 302b2b5 237774d d57909e 237774d d57909e 237774d 302b2b5 237774d 302b2b5 237774d 302b2b5 237774d d57909e 302b2b5 237774d |
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 |
import torch
import torch.nn as nn
from torchvision import models
class Stem(nn.Module):
def __init__(self):
super(Stem, self).__init__()
self.conv = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=3, stride=2),
nn.MaxPool2d(kernel_size=3, stride=2),
)
def forward(self, x):
x = self.conv(x)
return x
class ResidualBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride=1):
super().__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(
in_channels,
out_channels,
kernel_size=3,
stride=stride,
padding=1,
bias=False,
),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(inplace=True),
)
self.conv2 = nn.Sequential(
nn.Conv2d(
out_channels,
out_channels,
kernel_size=3,
stride=1,
padding=1,
bias=False,
),
nn.BatchNorm2d(out_channels),
)
self.shortcut = (
nn.Identity()
if in_channels == out_channels and stride == 1
else nn.Sequential(
nn.Conv2d(
in_channels, out_channels, kernel_size=1, stride=stride, bias=False
),
nn.BatchNorm2d(out_channels),
)
)
self.act = nn.LeakyReLU(inplace=True)
def forward(self, x):
identity = self.shortcut(x)
x = self.conv1(x)
x = self.conv2(x)
x += identity
return self.act(x)
class FromZero(nn.Module):
def __init__(self, num_classes=10):
super(FromZero, self).__init__()
self.stem = nn.Sequential(Stem())
self.layer1 = nn.Sequential(ResidualBlock(64, 64), ResidualBlock(64, 64))
self.layer2 = nn.Sequential(
ResidualBlock(64, 128, stride=2), ResidualBlock(128, 128)
)
self.layer3 = nn.Sequential(
ResidualBlock(128, 256, stride=2), ResidualBlock(256, 256)
)
self.layer4 = nn.Sequential(
ResidualBlock(256, 512, stride=2), ResidualBlock(512, 512), nn.Dropout(0.2)
)
self.flatten = nn.Flatten()
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Sequential(
nn.Linear(512, num_classes),
)
def forward(self, x):
x = self.stem(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = self.flatten(x)
x = self.fc(x)
return x
|