您当前的位置: 首页 > 

程序员正茂

暂无认证

  • 2浏览

    0关注

    283博文

    0收益

  • 0浏览

    0点赞

    0打赏

    0留言

私信
关注
热门博文

Win10下ImageAI-gpu训练自己的数据集

程序员正茂 发布时间:2020-03-27 21:31:07 ,浏览量:2

1.先按《Win10下ImageAI-gpu(tensorflow-gpu)环境安装(cuda及cuDNN自动安装)及实时对象检测》安装好。

2.下载pretrained-yolov3.h5,放到temp.py同一目录下。

3.下载hololens数据集,放到temp.py同一目录下。

4.编写代码开始训练。

from imageai.Detection.Custom import DetectionModelTrainer
trainer = DetectionModelTrainer()
trainer.setModelTypeAsYOLOv3()
trainer.setDataDirectory(data_directory="hololens")
trainer.setTrainConfig(object_names_array=["hololens"], batch_size=2, num_experiments=100, train_from_pretrained_model="pretrained-yolov3.h5")
trainer.trainModel()

batch_size=2是关键,否则可能出现以下错误

Resource exhausted: OOM when allocating tensor with shape[8,1024,14,14] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc

关注
打赏
1660743125
查看更多评论
立即登录/注册

微信扫码登录

0.0357s