def make_model():
set_seeds()
model = SeaNet({
1: (mm.MorphConv2d(1, 32, kernel_size=3, padding=1), 0),
2: (mm.MorphConv2d(32, 32, kernel_size=3, padding=1), 1),
3: (mm.MorphBCRLayer(32, 32, kernel_size=3, padding=1), 2),
4: (nn.MaxPool2d(2), 3),
5: (mm.MorphFlatLinear(6272, 10), 4)
})
model.eval()
return model
def make_seablock(in_planes, planes, stride=1, input_dim=32):
if stride == 1:
# simple
return SeaNet(
{
1: (mm.MorphConv2d(in_planes,
planes,
kernel_size=3,
padding=1,
stride=1,
bias=False), 0),
2: (mm.AddLayer(alpha=0.5), [1, 0]),
3: (mm.IdentityLayer(),
2), # To allow layers to be added at the end
},
input_shape=(in_planes, input_dim, input_dim),
tags='simple')
else:
# downsample
return SeaNet(
{
1: (mm.MorphConv2d(in_planes,
planes,
kernel_size=3,
padding=1,
stride=stride,
bias=False), 0),
2: (mm.MorphConv2d(in_planes,
planes,
kernel_size=1,
stride=stride,
bias=False), 0),
3: (mm.AddLayer(alpha=0.5), [2, 1]),
4: (mm.IdentityLayer(),
3), # To allow layers to be added at the end
},
input_shape=(in_planes, input_dim, input_dim),
tags='downsample')
set_seeds(args.seed)
for d in ['models', 'logs']:
try:
d_run = os.path.join('results', args.run_name, d)
assert not os.path.exists(d_run), '%s already exists!' % d_run
_ = os.makedirs(d_run)
except:
pass
base_model = SeaNet(
{
1: (mm.MorphBCRLayer(3, 64, kernel_size=3, padding=1), 0),
2: (mm.MaxPool2d(2), 1),
3: (mm.MorphBCRLayer(64, 128, kernel_size=3, padding=1), 2),
4: (mm.MaxPool2d(2), 3),
5: (mm.MorphBCRLayer(128, 128, kernel_size=3, padding=1), 4),
6: (mm.MorphFlatLinear(2048 * 4, 10), 5),
},
input_shape=(3, 32, 32))
pretrain_epochs = 20
additional_epochs = 20
dataloaders = make_dataloaders(train_size=1.0)
base_model, base_performance = train(base_model,
dataloaders,
epochs=pretrain_epochs,
verbose=True)
for p in base_performance: