- 以yolov3-tiny为例,加载函数是
ting.py
的load_darknet
- base.py
import torch
import torch.nn as nn
class BaseConv(nn.Module):
def __init__(self, in_channel, out_channel, k_size, stride, pad, bn, act):
super().__init__()
self.conv = nn.Conv2d(in_channel, out_channel, k_size, stride, pad, bias=not bn)
if bn == True:
self.bn = nn.BatchNorm2d(out_channel)
else:
self.bn = nn.Identity()
if act == 'leaky': self.act = nn.LeakyReLU(0.1)
elif act == 'linear': self.act = nn.Identity()
def forward(self, x):
x = self.conv(x)
x = self.bn(x)
return self.act(x)
if __name__ == '__main__':
data=torch.zeros(2,3, 320,320)
bc=BaseConv(3, 16, 3, 1, 1, False, 'leaky')
x = bc(data)
print(x.shape)
import torch
import torch.nn as nn
from base import BaseConv
import numpy as np
class Backbone(nn.Module):
def __init__(self):
super().__init__()
self.conv0 = BaseConv(3, 16, 3, 1, 1, True, 'leaky')
self.maxpool1 = nn.MaxPool2d(2)
self.conv2 = BaseConv(16, 32, 3, 1, 1, True, 'leaky')
self.maxpool3 = nn.MaxPool2d(2)
self.conv4 = BaseConv(32, 64, 3, 1, 1, True, 'leaky')
self.maxpool5 = nn.MaxPool2d(2)
self.conv6 = BaseConv(64, 128, 3, 1, 1, True, 'leaky')
self.maxpool7 = nn.MaxPool2d(2)
self.conv8 = BaseConv(128, 256, 3, 1, 1, True, 'leaky')
self.maxpool9 = nn.MaxPool2d(2)
self.conv10 = BaseConv(256, 512, 3, 1, 1, True, 'leaky')
self.pad = nn.ZeroPad2d((0,1,0,1))
self.maxpool11 = nn.MaxPool2d(2, 1)
self.conv12 = BaseConv(512, 1024, 3, 1, 1, True, 'leaky')
self.conv13 = BaseConv(1024, 256, 1, 1, 0, True, 'leaky')
self.conv14 = BaseConv(256, 512, 3, 1, 1, True, 'leaky')
def forward(self, x):
x0=self.conv0(x)
x1=self.maxpool1(x0)
x2=self.conv2(x1)
x3=self.maxpool3(x2)
x4=self.conv4(x3)
x5=self.maxpool5(x4)
x6=self.conv6(x5