直接上代码，源代码来自于kaggle上的大神，经过简单修改得到 # 模型 # 数据 # 训练 from torch. utils. dataimport Dataset, DataLoaderimport numpyfrom PILimport Imageclass convertDataset ( Dataset) : def __init__ ( self, data, transform) - > None : super ( ) . __init__( ) self. data= data self. transform= transform self. labels= self. data. iloc[ : , 0 ] . values self. images= self. data. iloc[ : , 1 : ] . values. astype( numpy. uint8) def __len__ ( self) : return len ( self. data) def __getitem__ ( self, index) : label= self. labels[ index] image= self. images[ index] . reshape( 28 , 28 ) image= Image. fromarray( image, "L" ) image= self. transform( image) return image, labeldef main ( ) : # 模型 import torch. nnas nnimport timm dict_path= "D:\\kaggle\\digit-recognizer\\dict.pth" import osclass VIT ( nn. Module) : def __init__ ( self, * args, ** kwargs) - > None : super ( ) . __init__( * args, ** kwargs) self. model= timm. create_model( "vit_base_patch16_224" , pretrained= True ) self. model. head= nn. Linear( self. model. head. in_features, 10 ) # type:ignore if os. path. exists( dict_path) : print ( "load dict" ) self. model. load_state_dict( torch. load( dict_path) ) def forward ( self, x) : return self. model( x) import torch device= torch. device( "cuda" if torch. cuda. is_availableelse "cpu" ) model= VIT( ) . to( device) # 数据 import pandas train_data= pandas. read_csv( "input/train.csv" ) from torchvisionimport transforms transform= transforms. Compose( [ transforms. Resize( ( 224 , 224 ) ) , transforms. Grayscale( num_output_channels= 3 ) , transforms. ToTensor( ) , transforms. Normalize( ( 0.5 , ) , ( 0.5 , ) ) ] ) train_dataset= convertDataset( train_data, transform) train_dataloader= DataLoader( train_dataset, batch_size= 20 , shuffle= True , num_workers= 1 , pin_memory= True ) # 训练 epoches= 5 optimizer= torch. optim. Adam( model. parameters( ) , lr= 1e - 4 ) criterion= nn. CrossEntropyLoss( ) import timeimport torch. autograd. profileras profilerfor epochin range ( epoches) : model. train( ) batch= 0 for images, labelsin train_dataloader: start_time= time. time( ) images, labels= images. to( device) , labels. to( device) optimizer. zero_grad( ) output= model( images) loss= criterion( output, labels) loss. backward( ) optimizer. step( ) batch+= 1 total= len ( labels) start_time1= time. time( ) _, predicted= output. max ( 1 ) correct= ( predicted== labels) . sum ( ) . item( ) print ( f"batch: { batch} / { len ( train_dataloader) } " , f"epoch: { epoch} / { epoches} " , f"loss: { loss. item( ) : .4f } " , f"accuracy: { correct/ total: .4f } " , f"spend: { time. time( ) - start_time: .2f } " ) if batch% 1000 == 0 : torch. save( model. model. state_dict( ) , dict_path) if __name__== "__main__" : main( ) 
