|
- from __future__ import print_function
- caffe_root = '/home/toanhoi/caffe/'
- import sys
- sys.path.insert(0, caffe_root + 'python')
- import caffe
- import math
- from caffe import layers as L
- from caffe.proto import caffe_pb2
-
- def bn_relu_conv_bn_relu(bottom, nout, dropout,split):
-
- if split == 'train':
- use_global_stats = False
- else:
- use_global_stats=True
-
- batch_norm1 = L.BatchNorm(bottom, batch_norm_param=dict(use_global_stats=use_global_stats), in_place=False,
- param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),
- dict(lr_mult=0, decay_mult=0)])
- scale1 = L.Scale(batch_norm1, bias_term=True, in_place=True,filler=dict(value=1), bias_filler=dict(value=0))
- relu1 = L.ReLU(scale1, in_place=True)
- conv1 = L.Convolution(relu1, kernel_size=[1, 1, 1], pad=[0, 0, 0], stride=[1,1,1],
- param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,
- num_output=nout * 4, axis=1, weight_filler=dict(type='msra'),
- bias_filler=dict(type='constant'))
-
- batch_norm2 = L.BatchNorm(conv1, batch_norm_param=dict(use_global_stats=use_global_stats), in_place=False,
- param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),
- dict(lr_mult=0, decay_mult=0)])
- scale2 = L.Scale(batch_norm2, bias_term=True, in_place=True,filler=dict(value=1), bias_filler=dict(value=0))
- relu2 = L.ReLU(scale2, in_place=True)
- conv2 = L.Convolution(relu2, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,
- axis=1, num_output=nout, pad=[1, 1, 1], kernel_size=[3, 3, 3], stride=[1,1,1],
- weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))
-
- if dropout > 0:
- conv2 = L.Dropout(conv2, dropout_ratio=dropout)
- return conv2
-
-
- def add_layer(bottom, num_filter, dropout,split):
- conv = bn_relu_conv_bn_relu(bottom, nout=num_filter, dropout=dropout,split=split)
- concate = L.Concat(bottom, conv, axis=1)
- return concate
-
-
- def transition(bottom, num_filter, split):
-
- if split == 'train':
- use_global_stats = False
- else:
- use_global_stats=True
-
- batch_norm1 = L.BatchNorm(bottom, batch_norm_param=dict(use_global_stats=use_global_stats), in_place=False,
- param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),
- dict(lr_mult=0, decay_mult=0)])
- scale1 = L.Scale(batch_norm1, bias_term=True, in_place=True,filler=dict(value=1), bias_filler=dict(value=0))
- relu1 = L.ReLU(scale1, in_place=True)
- conv1 = L.Convolution(relu1, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,
- axis=1, num_output=num_filter, pad=[0, 0, 0], kernel_size=[1, 1, 1],stride=[1,1,1],
- weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))
-
- batch_norm2 = L.BatchNorm(conv1, batch_norm_param=dict(use_global_stats=use_global_stats), in_place=False,
- param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),
- dict(lr_mult=0, decay_mult=0)])
- scale2 = L.Scale(batch_norm2, bias_term=True, in_place=True, filler=dict(value=1), bias_filler=dict(value=0))
- relu2 = L.ReLU(scale2, in_place=True)
-
- conv_down = L.Convolution(relu2, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,
- axis=1, num_output=num_filter, pad=[0, 0, 0], kernel_size=[2, 2, 2], stride=2,
- weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))
-
-
- #pooling = L.Pooling(conv1, type="Pooling", pool=P.Pooling.MAX, kernel_size=2, stride=2, engine=1)
- return conv_down
-
- # first_output -- #channels before entering the first dense block, set it to be comparable to growth_rate
- # growth_rate -- growth rate
- # dropout -- set to 0 to disable dropout, non-zero number to set dropout rate
- def densenet(split, batch_size=4, first_output=32, growth_rate=16, dropout=0.2):
- source_train_path = './train_list.txt'
- source_test_path = './test_list.txt'
- patch_size = [64, 64, 64]
- n = caffe.NetSpec()
- num_classes = 4
- reduction= 0.5
- N=[4,4,4,4]
- if split == 'train':
- n.data, n.label = L.HDF5Data(name="data", batch_size=batch_size, source=source_train_path, ntop=2, shuffle=True,
- transform_param=dict(crop_size_l=patch_size[0], crop_size_h=patch_size[1],
- crop_size_w=patch_size[2]), include={'phase': caffe.TRAIN})
- elif split == 'val':
- n.data, n.label = L.HDF5Data(name="data", batch_size=batch_size, source=source_test_path, ntop=2, shuffle=True,
- transform_param=dict(crop_size_l=patch_size[0], crop_size_h=patch_size[1],
- crop_size_w=patch_size[2]),
- include={'phase': caffe.TEST})
- else:
- n.data = L.Input(name="data", ntop=1, input_param={'shape': {'dim': [1, 2, patch_size[0], patch_size[1], patch_size[2]]}})
-
- nchannels = first_output
-
- # Fist layers
- n.conv1a = L.Convolution(n.data, param=[dict(lr_mult=1, decay_mult=1), dict(lr_mult=2, decay_mult=0)],
- axis=1, num_output=nchannels, pad=[1,1,1], kernel_size=[3, 3, 3], stride=[1,1,1],
- weight_filler=dict(type='msra'), bias_filler=dict(type='constant',value=-0.1))
-
- if split == 'train':
- use_global_stats = False
- else:
- use_global_stats=True
-
- n.bnorm1a = L.BatchNorm(n.conv1a, batch_norm_param=dict(use_global_stats=use_global_stats), param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),
- dict(lr_mult=0, decay_mult=0)], in_place=False)
-
- n.scale1a = L.Scale(n.bnorm1a, in_place=True, bias_term=True,filler=dict(value=1), bias_filler=dict(value=0))
- n.relu1a = L.ReLU(n.bnorm1a, in_place=True)
-
-
- # conv 1b, after BN set bias_term=false
- n.conv1b = L.Convolution(n.relu1a, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,
- axis=1, num_output=nchannels, pad=[1, 1, 1], kernel_size=[3, 3, 3], stride=[1,1,1],
- weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))
-
- n.bnorm1b = L.BatchNorm(n.conv1b, batch_norm_param=dict(use_global_stats=use_global_stats), param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),
- dict(lr_mult=0, decay_mult=0)], in_place=False)
- n.scale1b = L.Scale(n.bnorm1b, in_place=True, bias_term=True, filler=dict(value=1), bias_filler=dict(value=0))
- n.relu1b = L.ReLU(n.bnorm1b, in_place=True)
-
- n.conv1c = L.Convolution(n.relu1b, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,
- axis=1, num_output=nchannels, pad=[1, 1, 1], kernel_size=[3, 3, 3],stride=[1,1,1],
- weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))
- print (nchannels)
-
- # model = L.Pooling(n.conv1c, type="Pooling", pool=P.Pooling.MAX, kernel_size=2, stride=2, engine=1)
-
- n.bnorm1c = L.BatchNorm(n.conv1c, batch_norm_param=dict(use_global_stats=use_global_stats), param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),
- dict(lr_mult=0, decay_mult=0)], in_place=False)
- n.scale1c = L.Scale(n.bnorm1c, in_place=True, bias_term=True, filler=dict(value=1), bias_filler=dict(value=0))
- n.relu1c = L.ReLU(n.bnorm1c, in_place=True)
-
- model = L.Convolution(n.relu1c, param=[dict(lr_mult=1, decay_mult=1)], bias_term=False,
- axis=1, num_output=nchannels, pad=[0, 0, 0], kernel_size=[2, 2, 2], stride=2,
- weight_filler=dict(type='msra'), bias_filler=dict(type='constant'))
- n.__setattr__("Conv_down_1", model)
-
- # ===============Dense block 2=====================
- for i in range(N[0]):
- if (i == 0):
- concat = add_layer(model, growth_rate, dropout,split)
- n.__setattr__("Concat_%d" % (i + 1), concat)
- nchannels += growth_rate
- continue
- concat = add_layer(concat, growth_rate, dropout,split)
- n.__setattr__("Concat_%d" % (i + 1), concat)
- nchannels += growth_rate
- # ===============End dense block 2=================
- print (nchannels)
- # ===============Deconvolution layer 2==============
- model_deconv_x2 = L.Deconvolution(concat, param=[dict(lr_mult=0.1, decay_mult=1)],
- convolution_param=dict(kernel_size=[4,4,4], stride=[2,2,2], num_output=num_classes,
- pad=[1, 1, 1], group=num_classes,
- weight_filler=dict(type='bilinear_3D'),
- bias_term=False))
- n.__setattr__("Deconvolution_%d" % (N[0] + 1), model_deconv_x2)
- # ===============End Deconvolution layer 2==============
-
- # ===============Transition layer 2=================
- model = transition(concat, int(math.floor(nchannels * reduction)), split)
- n.__setattr__("Conv_down_%d" % (N[0] + 1), model)
- nchannels = int(math.floor(nchannels * reduction))
- # ===============End Transition layer2==============
-
- # ===============Dense block 3=====================
- for i in range(N[1]):
- if (i == 0):
- concat = add_layer(model, growth_rate, dropout, split)
- n.__setattr__("Concat_%d" % (N[1] + i + 2), concat)
- nchannels += growth_rate
- continue
- concat = add_layer(concat, growth_rate, dropout, split)
- n.__setattr__("Concat_%d" % (N[1] + i + 2), concat)
- nchannels += growth_rate
- # ===============End dense block 3=================
- print (nchannels)
- # ===============Deconvolution layer 3==============
- model_deconv_x4 = L.Deconvolution(concat, param=[dict(lr_mult=0.1, decay_mult=1)],
- convolution_param=dict(kernel_size=[6,6,6], stride=[4,4,4], num_output=num_classes,
- pad=[1, 1, 1], group=num_classes,
- weight_filler=dict(type='bilinear_3D'),
- bias_term=False))
- n.__setattr__("Deconvolution_%d" % (N[0] + N[1] + 2), model_deconv_x4)
- # ==============Transition layer 3=================
- model = transition(concat, int(math.floor(nchannels * reduction)), split)
- n.__setattr__("Conv_down_%d" % (N[0] + N[1] + 2), model)
- # ===============End Transition layer3==============
- nchannels = int(math.floor(nchannels * reduction))
-
- # ===============Dense block 4=====================
- for i in range(N[2]):
- if (i == 0):
- concat = add_layer(model, growth_rate, dropout, split)
- n.__setattr__("Concat_%d" % (N[0] + N[1] + i + 3), concat)
- nchannels += growth_rate
- continue
- concat = add_layer(concat, growth_rate, dropout, split)
- n.__setattr__("Concat_%d" % (N[0] + N[1] + i + 3), concat)
- nchannels += growth_rate
- # ===============End dense block 4=================
-
- # ===============Transition layer 4=================
- print(nchannels)
-
- # ===============Deconvolution layer 4==============
- model_deconv_x8 = L.Deconvolution(concat, param=[dict(lr_mult=0.1, decay_mult=1)],
- convolution_param=dict(kernel_size=[10,10,10], stride=[8,8,8], num_output=num_classes,
- pad=[1, 1, 1], group=num_classes,
- weight_filler=dict(type='bilinear_3D'),
- bias_term=False))
- n.__setattr__("Deconvolution_%d" % (N[0] + N[1] + N[2] + 3), model_deconv_x8)
- # ===============End Deconvolution layer 4==============
-
- # ===============Transition layer 4=================
- model = transition(concat, int(math.floor(nchannels * reduction)), split)
- n.__setattr__("Conv_down_%d" % (N[0] + N[1] + N[2] + 3), model)
- nchannels = int(math.floor(nchannels * reduction))
- # ===============End Transition layer3==============
-
- # ===============Dense block 5=====================
- for i in range(N[3]):
- if (i == 0):
- concat = add_layer(model, growth_rate, dropout, split)
- n.__setattr__("Concat_%d" % (N[0] + N[1] + N[2] + N[3] + i + 3), concat)
- nchannels += growth_rate
- continue
- concat = add_layer(concat, growth_rate, dropout, split)
- n.__setattr__("Concat_%d" % (N[0] + N[1] + N[2] + N[3] + i + 3), concat)
- nchannels += growth_rate
- # ===============End dense block 5=================
- print(nchannels)
-
- # ===============Deconvolution layer 5==============
- model_deconv_x16 = L.Deconvolution(concat, param=[dict(lr_mult=0.1, decay_mult=1)],
- convolution_param=dict(kernel_size=[18, 18, 18], stride=[16,16,16],
- num_output=num_classes,
- pad=[1, 1, 1], group=num_classes,
- weight_filler=dict(type='bilinear_3D'),
- bias_term=False))
- n.__setattr__("Deconvolution_%d" % (N[0] + N[1] + N[2] + N[3] + 4), model_deconv_x16)
- model = L.Concat(n.conv1c,model_deconv_x2, model_deconv_x4, model_deconv_x8, model_deconv_x16,
- axis=1)
-
- n.bnorm_concat= L.BatchNorm(model, batch_norm_param=dict(use_global_stats=use_global_stats), param=[dict(lr_mult=0, decay_mult=0), dict(lr_mult=0, decay_mult=0),
- dict(lr_mult=0, decay_mult=0)], in_place=False)
- n.scale_concat = L.Scale(n.bnorm_concat, in_place=True, bias_term=True, filler=dict(value=1), bias_filler=dict(value=0))
- n.relu_concat = L.ReLU(n.scale_concat, in_place=True)
- model_conv_concate = L.Convolution(n.relu_concat,
- param=[dict(lr_mult=1, decay_mult=1), dict(lr_mult=2, decay_mult=0)],
- axis=1, num_output=num_classes, pad=[0, 0, 0], kernel_size=[1, 1, 1],
- weight_filler=dict(type='msra'))
-
- if (split == 'train'):
- n.loss = L.SoftmaxWithLoss(model_conv_concate, n.label)
-
- elif (split == 'val'):
- n.loss = L.SoftmaxWithLoss(model_conv_concate, n.label)
- else:
- n.softmax = L.Softmax(model_conv_concate, ntop=1, in_place=False)
- return n.to_proto()
-
-
- def make_net():
- with open('train_3d_denseseg.prototxt', 'w') as f:
- print(str(densenet('train', batch_size=4)), file=f)
-
- with open('test_3d_denseseg.prototxt', 'w') as f:
- print(str(densenet('val', batch_size=4)), file=f)
- with open('deploy_3d_denseseg.prototxt', 'w') as f:
- print(str(densenet('deploy', batch_size=0)), file=f)
-
- def make_solver():
- s = caffe_pb2.SolverParameter()
- s.random_seed = 0xCAFFE
-
- s.train_net = 'train_3d_denseseg.prototxt'
-
- s.max_iter = 200000
- s.type = 'Adam'
- s.display = 20
-
- s.base_lr = 0.0002
- #s.power=0.9
-
- s.momentum = 0.97
- s.weight_decay = 0.0005
- s.average_loss=20
- s.iter_size = 1
- s.lr_policy='step'
- s.stepsize=50000
- s.gamma = 0.1
- s.snapshot_prefix ='./snapshot/3d_denseseg_iseg'
- s.snapshot = 2000
- s.solver_mode = caffe_pb2.SolverParameter.GPU
-
- solver_path = 'solver.prototxt'
- with open(solver_path, 'w') as f:
- f.write(str(s))
-
- if __name__ == '__main__':
- make_net()
- make_solver()
-
-
-
-
-
-
-
-
|