solve.py 627 字节
Newer Older
E
Evan Shelhamer 已提交
1
import caffe
2
import surgery, score
E
Evan Shelhamer 已提交
3 4 5

import numpy as np
import os
E
Evan Shelhamer 已提交
6
import sys
E
Evan Shelhamer 已提交
7 8 9 10

import setproctitle
setproctitle.setproctitle(os.path.basename(os.getcwd()))

11
weights = '../voc-fcn32s/voc-fcn32s.caffemodel'
E
Evan Shelhamer 已提交
12 13 14 15 16 17 18 19 20 21 22 23 24

# init
caffe.set_device(int(sys.argv[1]))
caffe.set_mode_gpu()

solver = caffe.SGDSolver('solver.prototxt')
solver.net.copy_from(weights)

# surgeries
interp_layers = [k for k in solver.net.params.keys() if 'up' in k]
surgery.interp(solver.net, interp_layers)

# scoring
E
Evan Shelhamer 已提交
25
val = np.loadtxt('../data/segvalid11.txt', dtype=str)
E
Evan Shelhamer 已提交
26 27 28 29

for _ in range(25):
    solver.step(4000)
    score.seg_tests(solver, False, val, layer='score')