parser.add_argument('--cols', dest='n', action='store', default=10000,
help='number of cols')
parser.add_argument('--r', dest='r', action='store', default=20,
help='internal dim')
parser.add_argument('--eps', dest='eps', action='store', default=1e-6,
help='ridge penalty')
parser.add_argument('--iter', dest='iter', action='store', default=10000,
help='max iter')
parser.add_argument('--set_from_master', dest='set_from_master', action='store_true',
help='samples are generated from the CPU of root: for obtaining identical dataset for different settings.')
args = parser.parse_args()
if args.with_gpu:
torch.cuda.set_device(rank % num_gpu)
if args.double:
TType=torch.cuda.DoubleTensor
else:
TType=torch.cuda.FloatTensor
else:
if args.double:
TType=torch.DoubleTensor
else:
TType=torch.FloatTensor
if args.nosubnormal:
torch.set_flush_denormal(True)
torch.manual_seed(95376+rank)
m = distmat.distgen_uniform(int(args.m), int(args.n), TType=TType, set_from_master=args.set_from_master)
nmf_driver = NMF(m, int(args.r), float(args.eps), TType, init_from_master=args.set_from_master)
nmf_driver.run(int(args.iter), tol=float(args.tol), check_interval=100, check_obj=True)
import torch
import torch.distributed as dist
dist.init_process_group('mpi')
rank = dist.get_rank()
size = dist.get_world_size()
device = 'cuda:{}'.format(rank) # or simply 'cpu' for CPU computing
if device.startswith('cuda'): torch.cuda.set_device(rank)
from dist_stat import distmat
A = distmat.distgen_uniform(4, 4, TType=torch.DoubleTensor)
B = distmat.distgen_uniform(4, 2, TType=torch.DoubleTensor)
AB = distmat.mm(A, B)
if rank == 0:
print("AB = ")
print(rank, AB.chunk)
C = (1 + AB).log()
if rank == 0:
print("log(1 + AB) = ")
print(rank, C.chunk)