isLandLZ commited on
Commit
796d9ba
1 Parent(s): 2767395

Upload model.py

Browse files
Files changed (1) hide show
  1. model.py +38 -0
model.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import jittor as jt
2
+ from jittor import nn, Module
3
+ import numpy as np
4
+ import sys, os
5
+ import random
6
+ import math
7
+ from jittor import init
8
+
9
+
10
+ class Model(Module):
11
+
12
+ def __init__(self):
13
+ super(Model, self).__init__()
14
+ # no padding
15
+ self.conv1 = nn.Conv(3, 32, 3, 1)
16
+ self.conv2 = nn.Conv(32, 64, 3, 1)
17
+ self.bn = nn.BatchNorm(64)
18
+
19
+ self.max_pool = nn.Pool(2, 2)
20
+ self.relu = nn.Relu()
21
+ self.fc1 = nn.Linear(64 * 12 * 12, 256)
22
+ self.fc2 = nn.Linear(256, 10)
23
+
24
+ def execute(self, x):
25
+ x = self.conv1(x)
26
+ x = self.relu(x)
27
+
28
+ x = self.conv2(x)
29
+ x = self.bn(x)
30
+ x = self.relu(x)
31
+
32
+ x = self.max_pool(x)
33
+ x = jt.reshape(x, [x.shape[0], -1])
34
+ x = self.fc1(x)
35
+ x = self.relu(x)
36
+ x = self.fc2(x)
37
+
38
+ return x