examples: imagenet: fix statistics computation

KingSpencer · Jul 5, 2016 · 998f9d9 · 998f9d9
1 parent f78b5e4
commit 998f9d9
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/examples/imagenet/cnn_imagenet.m b/examples/imagenet/cnn_imagenet.m
@@ -45,7 +45,10 @@
 else
   train = find(imdb.images.set == 1) ;
   images = fullfile(imdb.imageDir, imdb.images.name(train(1:100:end))) ;
-  [averageImage, rgbMean, rgbCovariance] = getImageStats(images, 'imageSize', [256 256]) ;
+  [averageImage, rgbMean, rgbCovariance] = getImageStats(images, ...
+                                                    'imageSize', [256 256], ...
+                                                    'numThreads', opts.numFetchThreads, ...
+                                                    'gpus', opts.train.gpus) ;
   save(imageStatsPath, 'averageImage', 'rgbMean', 'rgbCovariance') ;
 end
 [v,d] = eig(rgbCovariance) ;

diff --git a/examples/imagenet/getImageStats.m b/examples/imagenet/getImageStats.m
@@ -14,6 +14,7 @@
 numGpus = numel(opts.gpus) ;
 if numGpus > 0
   fprintf('%s: resetting GPU device\n', mfilename) ;
+  clear mex ;
   gpuDevice(opts.gpus(1))
 end
 
@@ -27,23 +28,23 @@
                        'imageSize', opts.imageSize, ...
                        'useGpu', numGpus > 0) ;
 
-  z = reshape(permute(data,[3 1 2 4]),3,[]) ;
-  n = size(z,2) ;
+  z = reshape(shiftdim(data,2),3,[]) ;
+  rgbm1{end+1} = mean(z,2) ;
+  rgbm2{end+1} = z*z'/size(z,2) ;
   avg{end+1} = mean(data, 4) ;
-  rgbm1{end+1} = sum(z,2)/n ;
-  rgbm2{end+1} = z*z'/n ;
   time = toc(time) ;
   fprintf(' %.1f Hz\n', numel(batch) / time) ;
 end
 
 averageImage = gather(mean(cat(4,avg{:}),4)) ;
-rgbm1 = mean(cat(2,rgbm1{:}),2) ;
-rgbm2 = mean(cat(3,rgbm2{:}),3) ;
-rgbMean = gather(rgbm1) ;
-rgbCovariance = gather(rgbm2 - rgbm1*rgbm1') ;
+rgbm1 = gather(mean(cat(2,rgbm1{:}),2)) ;
+rgbm2 = gather(mean(cat(3,rgbm2{:}),3)) ;
+rgbMean = rgbm1 ;
+rgbCovariance = rgbm2 - rgbm1*rgbm1' ;
 
 if numGpus > 0
   fprintf('%s: finished with GPU device, resetting again\n', mfilename) ;
+  clear mex ;
   gpuDevice(opts.gpus(1)) ;
 end
 fprintf('%s: all done\n', mfilename) ;