tensor_io.h

Go to the documentation of this file.

#ifndef MSHADOW_TENSOR_IO_H
#define MSHADOW_TENSOR_IO_H

#include <cstdio>
#include "tensor.h"

namespace mshadow{
    namespace utils{
        class IStream{
        public:
            virtual size_t Read( void *ptr, size_t size ) = 0;        
            virtual void Write( const void *ptr, size_t size ) = 0;
            virtual ~IStream( void ){}
        };
    };
    
    template<int dim,typename TStream>
    inline void SaveBinary( TStream &fo, const Tensor<cpu,dim> &src );
    template<int dim,typename TStream>
    inline void SaveBinary( TStream &fo, const Tensor<gpu,dim> &src );

    template<int dim,typename TStream>
    inline void LoadBinary( TStream &fi, Tensor<cpu,dim> &dst, bool pre_alloc );
    template<int dim,typename TStream>
    inline void LoadBinary( TStream &fi, Tensor<gpu,dim> &dst, bool pre_alloc );
    
    namespace utils{
        class FileStream: public IStream{
        public:
            FileStream( FILE *fp ):fp_(fp){}
            virtual size_t Read( void *ptr, size_t size ){
                return fread( ptr, size, 1, fp_ );
            }
            virtual void Write( const void *ptr, size_t size ){
                fwrite( ptr, size, 1, fp_ );
            }
            inline void Close( void ){
                fclose( fp_ );
            }
        private:
            FILE *fp_;
        };
    };
};

namespace mshadow{
    // implementations
    template<int dim, typename TStream>
    inline void SaveBinary( TStream &fo, const Tensor<cpu,dim> &src_ ){
        fo.Write( src_.shape.shape_, sizeof(index_t) * dim );
        Tensor<cpu,2> src = src_.FlatTo2D();
        for( index_t i = 0; i < src.shape[1]; ++ i ){
            fo.Write( src[i].dptr, sizeof(real_t)*src.shape[0] );
        }
    }
    template<int dim, typename TStream>
    inline void SaveBinary( TStream &fo, const Tensor<gpu,dim> &src ){
        // copy to CPU, then save
        Tensor<cpu,dim> tmp( src.shape ); 
        AllocSpace( tmp );
        Copy( tmp, src );
        SaveBinary( fo, tmp );
        FreeSpace( tmp );
    }

    template<int dim, typename TStream>
    inline void LoadBinary( TStream &fi, Tensor<cpu,dim> &dst_, bool pre_alloc ){
        Shape<dim> shape;
        utils::Assert( fi.Read( shape.shape_, sizeof(index_t) * dim ) != 0, "mshadow::LoadBinary" );
        if( pre_alloc ){
            utils::Assert( shape == dst_.shape );
        }else{
            dst_.shape = shape; AllocSpace( dst_ );
        }
        Tensor<cpu,2> dst = dst_.FlatTo2D();
        if( dst.shape[0] == 0 ) return;        
        for( index_t i = 0; i < dst.shape[1]; ++ i ){
            utils::Assert( fi.Read( dst[i].dptr, sizeof(real_t)*dst.shape[0] ) != 0, "mshadow::LoadBinary" );
        }
    } 
    template<int dim, typename TStream>
    inline void LoadBinary( TStream &fi, Tensor<gpu,dim> &dst, bool pre_alloc ){
        Tensor<cpu,dim> tmp;
        LoadBinary( fi, tmp, false );
        if( pre_alloc ){
            utils::Assert( tmp.shape == dst.shape );
        }else{
            dst.shape = tmp.shape; AllocSpace( dst );
        }
        Copy( dst, tmp );
        FreeSpace( tmp );
    }
};
#endif // TENSOR_IO_H