feng2022/gputest

1import numpy as np2import argparse3import functools4import os5import pickle6import sys7from datasets import Dataset8import gradio as gr9from pynvml import *10from transformers import pipeline11 12pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")13def predict(text):14  return pipe(text)[0]["translation_text"]15 16def print_gpu_utilization():17    nvmlInit()18    handle = nvmlDeviceGetHandleByIndex(0)19    info = nvmlDeviceGetMemoryInfo(handle)20    return f"GPU memory occupied: {info.used//1024**2} MB."21 22 23def print_summary(result):24    print(f"Time: {result.metrics['train_runtime']:.2f}")25    print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")26    print_gpu_utilization()27seq_len, dataset_size = 512, 51228dummy_data = {29    "input_ids": np.random.randint(100, 30000, (dataset_size, seq_len)),30    "labels": np.random.randint(0, 1, (dataset_size)),31}32ds = Dataset.from_dict(dummy_data)33ds.set_format("pt")34result = print_gpu_utilization()35iface = gr.Interface(36  fn=predict, 37  inputs='text',38  outputs='text',39  examples=[f'{result}']40)41 42iface.launch()