fcn-32s volumetric net

caffe_root = './caffe/'
import sys
sys.path.insert(0, caffe_root + 'python')
import caffe
from caffe import layers as L, params as P
from caffe.coord_map import crop

def conv_relu(bottom, nout, ks=3, stride=1, pad=1):
    conv = L.Convolution(bottom, kernel_size=ks, stride=stride, engine=1,
        num_output=nout, pad=pad, weight_filler=dict(type='xavier'),bias_filler=dict(type='constant', value=0)
        param=[dict(lr_mult=1, decay_mult=1), dict(lr_mult=2, decay_mult=0)])
    return conv, L.ReLU(conv, in_place=True)

def max_pool(bottom, ks=2, stride=2):
    return L.Pooling(bottom, pool=P.Pooling.AVE, kernel_size=ks, stride=stride,engine=2)

def fcn(split):
    n = caffe.NetSpec()
    batch_size=2
    source_train_path = './trainMS_list.txt'
    source_test_path = './testMS_list.txt'
    if split == 'train':
        n.data, n.label = L.HDF5Data(batch_size=batch_size, source=source_train_path, ntop=2,
                                     include={'phase': caffe.TRAIN})
    else:
        n.data, n.label = L.HDF5Data(batch_size=batch_size, source=source_test_path, ntop=2,
                                     include={'phase': caffe.TEST})

    # the base net
    n.conv1_1, n.relu1_1 = conv_relu(n.data, 32, pad=1)
    n.conv1_2, n.relu1_2 = conv_relu(n.relu1_1, 32)
    n.pool1 = max_pool(n.relu1_2)

    n.conv2_1, n.relu2_1 = conv_relu(n.pool1, 32)
    n.conv2_2, n.relu2_2 = conv_relu(n.relu2_1, 32)
    n.pool2 = max_pool(n.relu2_2)

    n.conv3_1, n.relu3_1 = conv_relu(n.pool2, 32)
    n.conv3_2, n.relu3_2 = conv_relu(n.relu3_1, 32)
    n.conv3_3, n.relu3_3 = conv_relu(n.relu3_2, 32)
    n.pool3 = max_pool(n.relu3_3)

    n.conv4_1, n.relu4_1 = conv_relu(n.pool3, 32)
    n.conv4_2, n.relu4_2 = conv_relu(n.relu4_1, 32)
    n.conv4_3, n.relu4_3 = conv_relu(n.relu4_2, 32)
    n.pool4 = max_pool(n.relu4_3)
    # Remove the conv 5 
    # fully conv
    n.fc6, n.relu6 = conv_relu(n.pool4, 32, ks=3, pad=0)
    n.drop6 = L.Dropout(n.relu6, dropout_ratio=0.5, in_place=True)
    n.fc7, n.relu7 = conv_relu(n.drop6, 32, ks=1, pad=0)
    n.drop7 = L.Dropout(n.relu7, dropout_ratio=0.5, in_place=True)
    n.score_fr = L.Convolution(n.drop7, num_output=4, kernel_size=1, pad=0,
                               param=[dict(lr_mult=1, decay_mult=1), dict(lr_mult=2, decay_mult=0)])
    n.upscore = L.Deconvolution(n.score_fr,
                                convolution_param=dict(num_output=4, kernel_size=20, stride=30,
                                                       bias_term=False),
                                param=[dict(lr_mult=0)])
    #n.score = crop(n.upscore, n.data)
    n.loss = L.SoftmaxWithLoss(n.upscore, n.label,
                               loss_param=dict(normalize=True))

    return n.to_proto()

def make_net():
    with open('train.prototxt', 'w') as f:
        f.write(str(fcn('train')))

    with open('val.prototxt', 'w') as f:
        f.write(str(fcn('val')))

if __name__ == '__main__':
    make_net()