| 63 | |
| 64 | |
| 65 | class Block(nn.Module): |
| 66 | def __init__(self, in_channels, out_channels, reps, strides=1, start_with_relu=True, grow_first=True): |
| 67 | super(Block, self).__init__() |
| 68 | |
| 69 | if out_channels != in_channels or strides != 1: |
| 70 | self.skip = nn.Conv2d(in_channels, out_channels, 1, stride=strides, bias=False) |
| 71 | self.skipbn = nn.BatchNorm2d(out_channels) |
| 72 | else: |
| 73 | self.skip = None |
| 74 | |
| 75 | rep = [] |
| 76 | for i in range(reps): |
| 77 | if grow_first: |
| 78 | inc = in_channels if i == 0 else out_channels |
| 79 | outc = out_channels |
| 80 | else: |
| 81 | inc = in_channels |
| 82 | outc = in_channels if i < (reps - 1) else out_channels |
| 83 | rep.append(nn.ReLU(inplace=True)) |
| 84 | rep.append(SeparableConv2d(inc, outc, 3, stride=1, padding=1)) |
| 85 | rep.append(nn.BatchNorm2d(outc)) |
| 86 | |
| 87 | if not start_with_relu: |
| 88 | rep = rep[1:] |
| 89 | else: |
| 90 | rep[0] = nn.ReLU(inplace=False) |
| 91 | |
| 92 | if strides != 1: |
| 93 | rep.append(nn.MaxPool2d(3, strides, 1)) |
| 94 | self.rep = nn.Sequential(*rep) |
| 95 | |
| 96 | def forward(self, inp): |
| 97 | x = self.rep(inp) |
| 98 | |
| 99 | if self.skip is not None: |
| 100 | skip = self.skip(inp) |
| 101 | skip = self.skipbn(skip) |
| 102 | else: |
| 103 | skip = inp |
| 104 | |
| 105 | x += skip |
| 106 | return x |
| 107 | |
| 108 | |
| 109 | class Xception(nn.Module): |