1.注意cuda与pytorch适配的问题
附上链接 卸载cuda的方法
注意安装对应的pytorch gpu版本
pytorch官网
2.有可能内存不够
下面是分内存的方法,当然你可以升级自己的硬件
内存分配
3.记得把xml文件转换为txt格式
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import glob2 as glob
classes = ["crazing", "inclusion", "patches", "pitted_surface", "rolled-in_scale", "scratches"]
def convert(size, box):
dw = 1./size[0]
dh = 1./size[1]
x = (box[0]+box[1])/2.0
y = (box[2]+box[3])/2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x, y, w, h)
def convert_anotation(image_name):
wd = os.getcwd()
in_flie = open('./ANNOTATIONS/'+image_name[:-3]+'xml')
out_flie = open('./LABELS/' + image_name[:-3] + 'txt', 'w')
tree = ET.parse(in_flie)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
clas = obj.find('name').text
if clas not in classes:
print(clas)
continue
clas_id = classes.index(clas)
xml_box = obj.find('bndbox')
b =(float(xml_box.find('xmin').text), float(xml_box.find('xmax').text),float(xml_box.find('ymin').text),float(xml_box.find('ymax').text))
bb = convert((w,h),b)
out_flie.write(str(clas_id)+" "+" ".join([str(a) for a in bb]) + '\n')
if __name__ == '__main__':
for image_path in glob.glob("./IMAGES/*.jpg"):
image_name = image_path.split('\\')[-1]
convert_anotation(image_name)