File size: 1,085 Bytes
fa791ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import torch
import time
import os
import argparse
import shutil
import sys
def parse_args():
parser = argparse.ArgumentParser(description='Matrix multiplication')
parser.add_argument('--gpus', help='gpu amount', required=True, type=int)
parser.add_argument('--size', help='matrix size', required=True, type=int)
parser.add_argument('--interval', help='sleep interval', required=True, type=float)
args = parser.parse_args()
return args
def matrix_multiplication(args):
a_list, b_list, result = [], [], []
size = (args.size, args.size)
# for i in range(args.gpus):
i = args.gpus
a_list.append(torch.rand(size, device=i))
b_list.append(torch.rand(size, device=i))
result.append(torch.rand(size, device=i))
while True:
# for i in range(args.gpus):
result[0] = a_list[0] * b_list[0]
time.sleep(args.interval)
if __name__ == "__main__":
# usage: python matrix_multiplication_gpus.py --size 20000 --gpus 2 --interval 0.01
args = parse_args()
matrix_multiplication(args)
|