batch_counter.py

from transformers import pipeline                                              
import random                                                                  
from torch.utils.data import Dataset                                           
import tqdm                                                                    
                                                                               
                                                                               
pipe = pipeline("text-classification", device=0)   
                                                                               
                                                                               
class MyDataset(Dataset):                                                      
    def __len__(self):                                                         
        return 1000                                                            
                                                                               
    def __getitem__(self, i):                                                  
        return "This is a test" * random.randrange(1, 30)                      
                                                                               
                                                                               
dataset = MyDataset()                                                          
                                                                               
print("-" * 30)                                                                
print("Streaming no batching")                                                 
                                                                               
for out in tqdm.tqdm(pipe(dataset)):                                           
    pass                                                                       
                                                                               
print("-" * 30)                                                                
print("Streaming batch_size=8")                                                
for out in tqdm.tqdm(pipe(dataset, batch_size=8), total=len(dataset)):         
    pass                                                                       
                                                                               
print("-" * 30)                                                                
print("Streaming batch_size=64")                                               
for out in tqdm.tqdm(pipe(dataset, batch_size=64), total=len(dataset)):        
    pass                                                                       
                                                                               
print("-" * 30)                                                                
print("Streaming batch_size=256")                                                                                                                                                                                                                                                         
for out in tqdm.tqdm(pipe(dataset, batch_size=256), total=len(dataset)):                                                                                                                                                                                                                  
    pass                                                                                                                                                                                                                                                                                  
                                                                                                                                                                                                                                                                                          
print("-" * 30)                                                                                                                                                                                                                                                                           
print("Streaming batch_size=512")                                                                                                                                                                                                                                                         
for out in tqdm.tqdm(pipe(dataset, batch_size=512), total=len(dataset)):                                                                                                                                                                                                                  
    pass                                                                                                                                                                                                                                                                                  
                                                                                                                                                                                                                                                                                          
print("-" * 30)                                                                                                                                                                                                                                                                           
print("Streaming batch_size=1024")                                                                                                                                                                                                                                                        
for out in tqdm.tqdm(pipe(dataset, batch_size=1024), total=len(dataset)):                                                                                                                                                                                                                 
    pass