深度学习应该选择tensorflow还是pytorch?一个场景分类任务惹出的问题

最近在做一个场景分类的任务,用AlexNet在Places数据集上做场景的分类训练,使用的框架是TF,自己不到2G的显存很容易就崩溃了,一直出现如下错误:

ResourceExhaustedError (see above for traceback): OOM when allocating tensor with shape[12544,4096]         
     [[Node: Variable_5/Adam_2/Assign = Assign[T=DT_FLOAT, _class=["loc:@Variable_5"], use_locking=true, validate_shape=true, _device="/job:localhost/replica:0/task:0/device:GPU:0"](Variable_5/Adam_2, Variable_5/Adam/Initializer/zeros)]]

不管是修改batchsize为1还是修改gpu选项都是报出同样的错误。

折腾了一个下午和晚上,还是选择换一个框架试试:pytorch

安装只需要一行命令:sudo pip3 install torch torchvision
(当然官网还可以根据操作系统与包的管理方式设置不同下载方式)

确实在安装上比tensorflow方便一点。

明天再用pytorch尝试周博磊博士的场景分类文章: Learning Deep Features for Scene Recognition using Places Database

猜你喜欢

转载自blog.csdn.net/weixin_39449466/article/details/80358220
今日推荐