You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
86 lines
3.0 KiB
86 lines
3.0 KiB
import torch |
|
import torch.nn as nn |
|
from torch.nn import functional as F |
|
|
|
|
|
class RestNetBasicBlock(nn.Module): |
|
def __init__(self, in_channels, out_channels, stride): |
|
super(RestNetBasicBlock, self).__init__() |
|
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1) |
|
self.bn1 = nn.BatchNorm2d(out_channels) |
|
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1) |
|
self.bn2 = nn.BatchNorm2d(out_channels) |
|
|
|
def forward(self, x): |
|
output = self.conv1(x) |
|
output = F.relu(self.bn1(output)) |
|
output = self.conv2(output) |
|
output = self.bn2(output) |
|
return F.relu(x + output) |
|
|
|
|
|
class RestNetDownBlock(nn.Module): |
|
def __init__(self, in_channels, out_channels, stride): |
|
super(RestNetDownBlock, self).__init__() |
|
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride[0], padding=1) |
|
self.bn1 = nn.BatchNorm2d(out_channels) |
|
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride[1], padding=1) |
|
self.bn2 = nn.BatchNorm2d(out_channels) |
|
self.extra = nn.Sequential( |
|
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride[0], padding=0), |
|
nn.BatchNorm2d(out_channels) |
|
) |
|
|
|
def forward(self, x): |
|
extra_x = self.extra(x) |
|
output = self.conv1(x) |
|
out = F.relu(self.bn1(output)) |
|
|
|
out = self.conv2(out) |
|
out = self.bn2(out) |
|
return F.relu(extra_x + out) |
|
|
|
|
|
class RestNet18(nn.Module): |
|
def __init__(self): |
|
super(RestNet18, self).__init__() |
|
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) |
|
self.bn1 = nn.BatchNorm2d(64) |
|
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) |
|
|
|
self.layer1 = nn.Sequential(RestNetBasicBlock(64, 64, 1), |
|
RestNetBasicBlock(64, 64, 1)) |
|
|
|
self.layer2 = nn.Sequential(RestNetDownBlock(64, 128, [2, 1]), |
|
RestNetBasicBlock(128, 128, 1)) |
|
|
|
self.layer3 = nn.Sequential(RestNetDownBlock(128, 256, [2, 1]), |
|
RestNetBasicBlock(256, 256, 1)) |
|
|
|
self.layer4 = nn.Sequential(RestNetDownBlock(256, 512, [2, 1]), |
|
RestNetBasicBlock(512, 512, 1)) |
|
|
|
self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1)) |
|
|
|
self.fc = nn.Linear(512, 10) |
|
|
|
def forward(self, x): |
|
out = self.conv1(x) |
|
out = self.layer1(out) |
|
out = self.layer2(out) |
|
out = self.layer3(out) |
|
out = self.layer4(out) |
|
out = self.avgpool(out) |
|
out = out.reshape(x.shape[0], -1) |
|
out = self.fc(out) |
|
return out |
|
|
|
|
|
if __name__ == '__main__': |
|
# build model |
|
model = RestNet18() |
|
model.eval() |
|
|
|
# export onnx (rknn-toolkit2 only support opset_version=12) |
|
x = torch.randn((1, 3, 224, 224)) |
|
torch.onnx.export(model, x, './resnet18.onnx', opset_version=12, input_names=['input'], output_names=['output'])
|
|
|