0

我正在尝试从预训练的 ResNet 模型中提取特征。

但是在某些时候它会因错误而崩溃：

F0131 01:10:06.501891 17766 math_functions.cu:79] Check failed: error == cudaSuccess (4 vs. 0)  unspecified launch failure
*** Check failure stack trace: ***

Howewer 我已经测试了 AlexNet 模型，看起来还可以。nvidia-smi显示 GPU 内存未满（我有 gtx 1070 和 8Gb 内存，批量大小为 1，所以应该不是问题）。

这种行为的原因是什么？

这是代码：

#Code based on http://nbviewer.jupyter.org/github/BVLC/caffe/blob/master/examples/00-classification.ipynb

import numpy as np
import sys
import os

os.environ['GLOG_minloglevel'] = '2' 
caffe_root = '/media/myuser/120Gb/caffe-master/caffe/'  # parameter
sys.path.insert(0, caffe_root + 'python')
import caffe

if len(sys.argv) != 5:
    print "Usage: python feature_extract.py deploy.prototxt model.caffemodel mean.npy <train_data_dir>"
    sys.exit()

DEPLOY_PROTOTXT= sys.argv[1] #deploy.prototxt
CAFFE_MODEL= sys.argv[2] #model.caffemodel
MEAN_FILE= sys.argv[3] #mean.npy
TRAIN_DATA_DIR= sys.argv[4]

MODEL_POSTFIX= CAFFE_MODEL.split(os.sep)[-2] # not safe way
print "MODEL_POSTFIX", MODEL_POSTFIX

LAYER_NAME='fc1000'
LAYER_DIM=1000

#caffe.set_mode_cpu()
caffe.set_mode_gpu()

net = caffe.Net(DEPLOY_PROTOTXT, CAFFE_MODEL, caffe.TEST)
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_transpose('data', (2,0,1))
transformer.set_mean('data', np.load(MEAN_FILE).mean(1).mean(1))
transformer.set_raw_scale('data', 255)
transformer.set_channel_swap('data', (2,1,0))

net.blobs['data'].reshape(1,3,224,224)
#net.blobs['data'].reshape(1,3,227,227)

#main
if LAYER_NAME not in net.blobs:
    raise TypeError("Invalid layer name: " + LAYER_NAME)

filename_list= []
label_list= []
d = {'cat': 0, 'dog': 1}
for key, value in d.iteritems():
    print TRAIN_DATA_DIR+"/"+key 
    for file in os.listdir(TRAIN_DATA_DIR+"/"+key):
        if file.endswith(".jpg"):
            filename_list.append(TRAIN_DATA_DIR+"/"+key+"/"+file)
            label_list.append(value)
print "total_number_of_images", len(filename_list)

features= np.zeros((len(filename_list), LAYER_DIM), dtype=np.float32) # Note: number of features is hardcoded.

for i in range(0, len(filename_list)):
    #if(i % 1000 == 0):
    print i,"/",len(filename_list)

    print "net.blobs['data'].data", net.blobs['data'].data.shape
    net.blobs['data'].data[0] = transformer.preprocess('data', caffe.io.load_image(filename_list[i]))

    net.forward()

    print "np.ravel(net.blobs[LAYER_NAME].data)", np.ravel(net.blobs[LAYER_NAME].data).shape

    features[i,:]= np.ravel(net.blobs[LAYER_NAME].data)

#Save features
features_filename= "kaggle_data/X.npy"
np.save(features_filename, features)
#Save labels
labels_filename= "kaggle_data/Y.npy"
np.save(labels_filename, np.array(label_list))
4

0 回答 0