def test_randperm(self): state = torch.random.get_rng_state() # results a = ht.random.randperm(10, dtype=ht.int32) b = ht.random.randperm(4, dtype=ht.float32, split=0) c = ht.random.randperm(5, split=0) d = ht.random.randperm(5, dtype=ht.float64) torch.random.set_rng_state(state) # torch results to compare to a_cmp = torch.randperm(10, dtype=torch.int32, device=self.device.torch_device) b_cmp = torch.randperm(4, dtype=torch.float32, device=self.device.torch_device) c_cmp = torch.randperm(5, dtype=torch.int64, device=self.device.torch_device) d_cmp = torch.randperm(5, dtype=torch.float64, device=self.device.torch_device) self.assertEqual(a.dtype, ht.int32) self.assertTrue((a.larray == a_cmp).all()) self.assertEqual(b.dtype, ht.float32) self.assertTrue((ht.resplit(b).larray == b_cmp).all()) self.assertEqual(c.dtype, ht.int64) self.assertTrue((ht.resplit(c).larray == c_cmp).all()) self.assertEqual(d.dtype, ht.float64) self.assertTrue((d.larray == d_cmp).all()) with self.assertRaises(TypeError): ht.random.randperm("abc")
def fit(self, X, y, sample_weight=None): """ Fit Gaussian Naive Bayes according to X, y Parameters ---------- X : ht.tensor of shape (n_samples, n_features) Training set, where n_samples is the number of samples and n_features is the number of features. y : ht.tensor of shape (n_samples,) Labels for training set. sample_weight : ht.tensor of shape (n_samples,), optional (default=None) Weights applied to individual samples (1. for unweighted). Returns ------- self : object """ # sanitize input - to be moved to sanitation module, cf. #468 if not isinstance(X, ht.DNDarray): raise ValueError( "input needs to be a ht.DNDarray, but was {}".format(type(X))) if not isinstance(y, ht.DNDarray): raise ValueError( "input needs to be a ht.DNDarray, but was {}".format(type(y))) if y.numdims != 1: raise ValueError("expected y to be a 1-D tensor, is {}-D".format( y.numdims)) if sample_weight is not None: if not isinstance(sample_weight, ht.DNDarray): raise ValueError( "sample_weight needs to be a ht.DNDarray, but was {}". format(type(sample_weight))) classes = ht.unique(y, sorted=True) if classes.split is not None: classes = ht.resplit(classes, axis=None) return self.__partial_fit(X, y, classes, _refit=True, sample_weight=sample_weight)
def test_misc_coverage(self): length = torch.tensor([i + 5 for i in range(3)], device=self.device.torch_device) test = torch.arange( torch.prod(length), dtype=torch.float64, device=self.device.torch_device ).reshape([i + 5 for i in range(3)]) a = ht.array(test, split=None) tiles = ht.tiling.SplitTiles(a) self.assertTrue(torch.all(tiles.tile_locations == a.comm.rank)) a = ht.resplit(a, 0) tiles = ht.tiling.SplitTiles(a) if a.comm.size == 3: # definition of adjusting tests is he same logic as the code itself, # therefore, fixed tests are issued for one process confic tile_dims = torch.tensor( [[2.0, 2.0, 1.0], [2.0, 2.0, 2.0], [3.0, 2.0, 2.0]], device=self.device.torch_device ) res = tiles.tile_dimensions self.assertTrue(torch.equal(tile_dims, res)) testing_tensor = torch.tensor( [ [ [168.0, 169.0, 170.0, 171.0, 172.0, 173.0, 174.0], [175.0, 176.0, 177.0, 178.0, 179.0, 180.0, 181.0], [182.0, 183.0, 184.0, 185.0, 186.0, 187.0, 188.0], [189.0, 190.0, 191.0, 192.0, 193.0, 194.0, 195.0], [196.0, 197.0, 198.0, 199.0, 200.0, 201.0, 202.0], [203.0, 204.0, 205.0, 206.0, 207.0, 208.0, 209.0], ] ], dtype=torch.float64, device=self.device.torch_device, ) if a.comm.rank == 2: self.assertTrue(torch.equal(tiles[2], testing_tensor)) tiles[2] = 1000 sl = tiles[2] if a.comm.rank == 2: self.assertEqual(torch.Size([1, 6, 7]), sl.shape) self.assertTrue(torch.all(sl == 1000)) else: self.assertTrue(sl is None)
def fit(self, x: DNDarray, y: DNDarray, sample_weight: Optional[DNDarray] = None): """ Fit Gaussian Naive Bayes according to ``x`` and ``y`` Parameters ---------- x : DNDarray Training set, where n_samples is the number of samples and n_features is the number of features. Shape = (n_classes, n_features) y : DNDarray Labels for training set. Shape = (n_samples, ) sample_weight : DNDarray, optional Weights applied to individual samples (1. for unweighted). Shape = (n_samples, ) """ # sanitize input - to be moved to sanitation module, cf. #468 if not isinstance(x, ht.DNDarray): raise ValueError( "input needs to be a ht.DNDarray, but was {}".format(type(x))) if not isinstance(y, ht.DNDarray): raise ValueError( "input needs to be a ht.DNDarray, but was {}".format(type(y))) if y.ndim != 1: raise ValueError("expected y to be a 1-D tensor, is {}-D".format( y.ndim)) if sample_weight is not None: if not isinstance(sample_weight, ht.DNDarray): raise ValueError( "sample_weight needs to be a ht.DNDarray, but was {}". format(type(sample_weight))) classes = ht.unique(y, sorted=True) if classes.split is not None: classes = ht.resplit(classes, axis=None) return self.__partial_fit(x, y, classes, _refit=True, sample_weight=sample_weight)
def test_permutation(self): # Reset RNG ht.random.seed() state = torch.random.get_rng_state() # results a = ht.random.permutation(10) b_arr = ht.arange(10, dtype=ht.float32) b = ht.random.permutation(ht.resplit(b_arr, 0)) c_arr = ht.arange(16).reshape((4, 4)) c = ht.random.permutation(c_arr) c0 = ht.random.permutation(ht.resplit(c_arr, 0)) c1 = ht.random.permutation(ht.resplit(c_arr, 1)) torch.set_rng_state(state) # torch results to compare to a_cmp = torch.randperm(a.shape[0], device=self.device.torch_device) b_cmp = b_arr.larray[torch.randperm(b.shape[0], device=self.device.torch_device)] c_cmp = c_arr.larray[torch.randperm(c.shape[0], device=self.device.torch_device)] c0_cmp = c_arr.larray[torch.randperm(c.shape[0], device=self.device.torch_device)] c1_cmp = c_arr.larray[torch.randperm(c.shape[0], device=self.device.torch_device)] # compare self.assertEqual(a.dtype, ht.int64) self.assertTrue((a.larray == a_cmp).all()) self.assertEqual(b.dtype, ht.float32) self.assertTrue((ht.resplit(b).larray == b_cmp).all()) self.assertTrue((c.larray == c_cmp).all()) self.assertTrue((ht.resplit(c0).larray == c0_cmp).all()) self.assertTrue((ht.resplit(c1).larray == c1_cmp).all()) with self.assertRaises(TypeError): ht.random.permutation("abc")
def test_resplit(self): # resplitting with same axis, should leave everything unchanged shape = (ht.MPI_WORLD.size, ht.MPI_WORLD.size) data = ht.zeros(shape, split=None, device=ht_device) data2 = ht.resplit(data, None) self.assertIsInstance(data2, ht.DNDarray) self.assertEqual(data2.shape, shape) self.assertEqual(data2.lshape, shape) self.assertEqual(data2.split, None) # resplitting with same axis, should leave everything unchanged shape = (ht.MPI_WORLD.size, ht.MPI_WORLD.size) data = ht.zeros(shape, split=1, device=ht_device) data2 = ht.resplit(data, 1) self.assertIsInstance(data2, ht.DNDarray) self.assertEqual(data2.shape, shape) self.assertEqual(data2.lshape, (data.comm.size, 1)) self.assertEqual(data2.split, 1) # splitting an unsplit tensor should result in slicing the tensor locally shape = (ht.MPI_WORLD.size, ht.MPI_WORLD.size) data = ht.zeros(shape, device=ht_device) data2 = ht.resplit(data, 1) self.assertIsInstance(data2, ht.DNDarray) self.assertEqual(data2.shape, shape) self.assertEqual(data2.lshape, (data.comm.size, 1)) self.assertEqual(data2.split, 1) # unsplitting, aka gathering a tensor shape = (ht.MPI_WORLD.size + 1, ht.MPI_WORLD.size) data = ht.ones(shape, split=0, device=ht_device) data2 = ht.resplit(data, None) self.assertIsInstance(data2, ht.DNDarray) self.assertEqual(data2.shape, shape) self.assertEqual(data2.lshape, shape) self.assertEqual(data2.split, None) # assign and entirely new split axis shape = (ht.MPI_WORLD.size + 2, ht.MPI_WORLD.size + 1) data = ht.ones(shape, split=0, device=ht_device) data2 = ht.resplit(data, 1) self.assertIsInstance(data2, ht.DNDarray) self.assertEqual(data2.shape, shape) self.assertEqual(data2.lshape[0], ht.MPI_WORLD.size + 2) self.assertTrue(data2.lshape[1] == 1 or data2.lshape[1] == 2) self.assertEqual(data2.split, 1) # test sorting order of resplit N = ht.MPI_WORLD.size reference_tensor = ht.zeros((N, N + 1, 2 * N)) for n in range(N): for m in range(N + 1): reference_tensor[n, m, :] = ht.arange(0, 2 * N) + m * 10 + n * 100 # split along axis = 0 resplit_tensor = ht.resplit(reference_tensor, axis=0) local_shape = (1, N + 1, 2 * N) local_tensor = reference_tensor[ht.MPI_WORLD.rank, :, :] self.assertEqual(resplit_tensor.lshape, local_shape) self.assertTrue((resplit_tensor._DNDarray__array == local_tensor._DNDarray__array).all()) # unsplit unsplit_tensor = ht.resplit(resplit_tensor, axis=None) self.assertTrue( (unsplit_tensor._DNDarray__array == reference_tensor._DNDarray__array).all() ) # split along axis = 1 resplit_tensor = ht.resplit(unsplit_tensor, axis=1) if ht.MPI_WORLD.rank == 0: local_shape = (N, 2, 2 * N) local_tensor = reference_tensor[:, 0:2, :] else: local_shape = (N, 1, 2 * N) local_tensor = reference_tensor[:, ht.MPI_WORLD.rank + 1 : ht.MPI_WORLD.rank + 2, :] self.assertEqual(resplit_tensor.lshape, local_shape) self.assertTrue((resplit_tensor._DNDarray__array == local_tensor._DNDarray__array).all()) # unsplit unsplit_tensor = ht.resplit(resplit_tensor, axis=None) self.assertTrue( (unsplit_tensor._DNDarray__array == reference_tensor._DNDarray__array).all() ) # split along axis = 2 resplit_tensor = ht.resplit(unsplit_tensor, axis=2) local_shape = (N, N + 1, 2) local_tensor = reference_tensor[:, :, 2 * ht.MPI_WORLD.rank : 2 * ht.MPI_WORLD.rank + 2] self.assertEqual(resplit_tensor.lshape, local_shape) self.assertTrue((resplit_tensor._DNDarray__array == local_tensor._DNDarray__array).all())
def test_fit_iris(self): # load sklearn train/test sets and resulting probabilities X_train = ht.load( "heat/datasets/data/iris_X_train.csv", sep=";", dtype=ht.float64, device=ht_device ) X_test = ht.load( "heat/datasets/data/iris_X_test.csv", sep=";", dtype=ht.float64, device=ht_device ) y_train = ht.load( "heat/datasets/data/iris_y_train.csv", sep=";", dtype=ht.int64, device=ht_device ).squeeze() y_test = ht.load( "heat/datasets/data/iris_y_test.csv", sep=";", dtype=ht.int64, device=ht_device ).squeeze() y_pred_proba_sklearn = ht.load( "heat/datasets/data/iris_y_pred_proba.csv", sep=";", dtype=ht.float64, device=ht_device ) # test ht.GaussianNB from heat.naive_bayes import GaussianNB gnb_heat = GaussianNB() self.assertEqual(gnb_heat.priors, None) with self.assertRaises(AttributeError): gnb_heat.classes_ with self.assertRaises(AttributeError): gnb_heat.class_prior_ with self.assertRaises(AttributeError): gnb_heat.epsilon_ # test GaussianNB locally, no weights local_fit = gnb_heat.fit(X_train, y_train) self.assert_array_equal(gnb_heat.classes_, np.array([0, 1, 2])) local_fit_no_classes = gnb_heat.partial_fit(X_train, y_train, classes=None) y_pred_local = local_fit_no_classes.predict(X_test) y_pred_proba_local = local_fit.predict_proba(X_test) sklearn_class_prior = np.array([0.38666667, 0.26666667, 0.34666667]) sklearn_epsilon = np.array([3.6399040000000003e-09]) sklearn_theta = ht.array( [ [4.97586207, 3.35862069, 1.44827586, 0.23448276], [5.935, 2.71, 4.185, 1.3], [6.77692308, 3.09230769, 5.73461538, 2.10769231], ], dtype=X_train.dtype, device=ht_device, ) sklearn_sigma = ht.array( [ [0.10321047, 0.13208086, 0.01629013, 0.00846612], [0.256275, 0.0829, 0.255275, 0.046], [0.38869823, 0.10147929, 0.31303255, 0.04763314], ], dtype=X_train.dtype, device=ht_device, ) self.assertIsInstance(y_pred_local, ht.DNDarray) self.assertEqual((y_pred_local != y_test).sum(), ht.array(4)) self.assert_array_equal(gnb_heat.class_prior_, sklearn_class_prior) self.assert_array_equal(gnb_heat.epsilon_, sklearn_epsilon) self.assertTrue(ht.isclose(gnb_heat.theta_, sklearn_theta).all()) self.assertTrue(ht.isclose(gnb_heat.sigma_, sklearn_sigma, atol=1e-1).all()) self.assertTrue(ht.isclose(y_pred_proba_sklearn, y_pred_proba_local, atol=1e-1).all()) # test GaussianNB when sample_weight is not None, sample_weight not distributed sample_weight = ht.ones((y_train.gshape[0]), dtype=ht.float32, split=None) local_fit_weight = gnb_heat.fit(X_train, y_train, sample_weight=sample_weight) y_pred_local_weight = local_fit_weight.predict(X_test) y_pred_proba_local_weight = local_fit_weight.predict_proba(X_test) self.assertIsInstance(y_pred_local_weight, ht.DNDarray) self.assert_array_equal(gnb_heat.class_prior_, sklearn_class_prior) self.assert_array_equal(gnb_heat.epsilon_, sklearn_epsilon) self.assertTrue(ht.isclose(gnb_heat.theta_, sklearn_theta).all()) self.assertTrue(ht.isclose(gnb_heat.sigma_, sklearn_sigma, atol=1e-1).all()) self.assert_array_equal(y_pred_local_weight, y_pred_local.numpy()) self.assertTrue(ht.isclose(y_pred_proba_sklearn, y_pred_proba_local_weight).all()) # test GaussianNB, data and labels distributed along split axis 0 X_train_split = ht.resplit(X_train, axis=0) X_test_split = ht.resplit(X_test, axis=0) y_train_split = ht.resplit(y_train, axis=0) y_test_split = ht.resplit(y_test, axis=0) y_pred_split = gnb_heat.fit(X_train_split, y_train_split).predict(X_test_split) self.assert_array_equal(gnb_heat.class_prior_, sklearn_class_prior) self.assert_array_equal(gnb_heat.epsilon_, sklearn_epsilon) self.assertTrue(ht.isclose(gnb_heat.theta_, sklearn_theta).all()) self.assertTrue(ht.isclose(gnb_heat.sigma_, sklearn_sigma, atol=1e-1).all()) self.assert_array_equal(y_pred_split, y_pred_local.numpy()) self.assertEqual((y_pred_split != y_test_split).sum(), ht.array(4)) sample_weight_split = ht.ones(y_train_split.gshape[0], dtype=ht.float32, split=0) y_pred_split_weight = gnb_heat.fit( X_train_split, y_train_split, sample_weight=sample_weight_split ).predict(X_test_split) self.assertIsInstance(y_pred_split_weight, ht.DNDarray) self.assert_array_equal(y_pred_split_weight, y_pred_split.numpy()) # test exceptions X_torch = torch.ones(75, 4) y_np = np.zeros(75) y_2D = ht.ones((75, 1), split=None, device=ht_device) weights_torch = torch.zeros(75) X_3D = ht.ones((75, 4, 4), split=None, device=ht_device) X_wrong_size = ht.ones((75, 5), split=None, device=ht_device) y_wrong_size = ht.zeros(76, device=ht_device) X_train_split = ht.resplit(X_train, axis=0) y_train_split = ht.resplit(y_train, axis=0) weights_2D_split = y_2D = ht.ones((75, 1), split=0, device=ht_device) weights_wrong_size = ht.ones(76, device=ht_device) priors_wrong_shape = ht.random.randn(4, device=ht_device) priors_wrong_sum = ht.random.randn(3, dtype=ht.float32, device=ht_device) priors_wrong_sign = ht.array([-0.3, 0.7, 0.6]) wrong_classes = ht.array([3, 4, 5]) with self.assertRaises(ValueError): gnb_heat.fit(X_torch, y_train) with self.assertRaises(ValueError): gnb_heat.fit(X_train, y_np) with self.assertRaises(ValueError): gnb_heat.fit(X_train, y_2D) with self.assertRaises(ValueError): gnb_heat.fit(X_train, y_train, sample_weight=weights_torch) with self.assertRaises(ValueError): gnb_heat.fit(X_3D, y_train) with self.assertRaises(ValueError): gnb_heat.fit(X_train, y_wrong_size) with self.assertRaises(ValueError): gnb_heat.fit(X_train, y_train) gnb_heat.predict(X_torch) with self.assertRaises(ValueError): gnb_heat.fit(X_train, y_train) gnb_heat.partial_fit(X_wrong_size, y_train) with self.assertRaises(ValueError): gnb_heat.fit(X_train, y_train) gnb_heat.partial_fit(X_train, y_train, classes=wrong_classes) with self.assertRaises(ValueError): gnb_heat.classes_ = None gnb_heat.partial_fit(X_train, y_train, classes=None) with self.assertRaises(ValueError): gnb_heat.fit(X_train_split, y_train_split, sample_weight=weights_2D_split) with self.assertRaises(ValueError): gnb_heat.fit(X_train, y_train, sample_weight=weights_wrong_size) with self.assertRaises(ValueError): gnb_heat.priors = priors_wrong_shape gnb_heat.fit(X_train, y_train) with self.assertRaises(ValueError): gnb_heat.priors = priors_wrong_sum gnb_heat.fit(X_train, y_train) with self.assertRaises(ValueError): gnb_heat.priors = priors_wrong_sign gnb_heat.fit(X_train, y_train)
def __partial_fit(self, X, y, classes=None, _refit=False, sample_weight=None): """ Actual implementation of Gaussian NB fitting. Adapted to HeAT from scikit-learn. Parameters ---------- X : ht.tensor of shape (n_samples, n_features) Training set, where n_samples is the number of samples and n_features is the number of features. y : ht.tensor of shape (n_samples,) Labels for training set. classes : ht.tensor of shape (n_classes,), optional (default=None) List of all the classes that can possibly appear in the y vector. Must be provided at the first call to partial_fit, can be omitted in subsequent calls. _refit : bool, optional (default=False) If true, act as though this were the first time __partial_fit is called (ie, throw away any past fitting and start over). sample_weight : ht.tensor of shape (n_samples,), optional (default=None) Weights applied to individual samples (1. for unweighted). Returns ------- self : object """ # TODO: sanitize X and y shape: sanitation/validation module, cf. #468 n_samples = X.shape[0] if X.numdims != 2: raise ValueError("expected X to be a 2-D tensor, is {}-D".format( X.numdims)) if y.shape[0] != n_samples: raise ValueError( "y.shape[0] must match number of samples {}, is {}".format( n_samples, y.shape[0])) # TODO: sanitize sample_weight: sanitation/validation module, cf. #468 if sample_weight is not None: if sample_weight.numdims != 1: raise ValueError("Sample weights must be 1D tensor") if sample_weight.shape != (n_samples, ): raise ValueError( "sample_weight.shape == {}, expected {}!".format( sample_weight.shape, (n_samples, ))) # If the ratio of data variance between dimensions is too small, it # will cause numerical errors. To address this, we artificially # boost the variance by epsilon, a small fraction of the standard # deviation of the largest dimension. self.epsilon_ = self.var_smoothing * ht.var(X, axis=0).max() if _refit: self.classes_ = None if self.__check_partial_fit_first_call(classes): # This is the first call to partial_fit: # initialize various cumulative counters n_features = X.shape[1] n_classes = len(self.classes_) self.theta_ = ht.zeros((n_classes, n_features), dtype=X.dtype, device=X.device) self.sigma_ = ht.zeros((n_classes, n_features), dtype=X.dtype, device=X.device) self.class_count_ = ht.zeros((n_classes, ), dtype=ht.float64, device=X.device) # Initialise the class prior # Take into account the priors if self.priors is not None: if not isinstance(self.priors, ht.DNDarray): priors = ht.array(self.priors, dtype=X.dtype, split=None, device=X.device) else: priors = self.priors # Check that the provide prior match the number of classes if len(priors) != n_classes: raise ValueError("Number of priors must match number of" " classes.") # Check that the sum is 1 if not ht.isclose(priors.sum(), ht.array(1.0, dtype=priors.dtype)): raise ValueError("The sum of the priors should be 1.") # Check that the prior are non-negative if (priors < 0).any(): raise ValueError("Priors must be non-negative.") self.class_prior_ = priors else: # Initialize the priors to zeros for each class self.class_prior_ = ht.zeros(len(self.classes_), dtype=ht.float64, split=None, device=X.device) else: if X.shape[1] != self.theta_.shape[1]: raise ValueError( "Number of features {} does not match previous data {}.". format(X.shape[1], self.theta_.shape[1])) # Put epsilon back in each time self.sigma_[:, :] -= self.epsilon_ classes = self.classes_ unique_y = ht.unique(y, sorted=True) if unique_y.split is not None: unique_y = ht.resplit(unique_y, axis=None) unique_y_in_classes = ht.eq(unique_y, classes) if not ht.all(unique_y_in_classes): raise ValueError("The target label(s) {} in y do not exist in the " "initial classes {}".format( unique_y[~unique_y_in_classes], classes)) for y_i in unique_y: # assuming classes.split is None if y_i in classes: i = ht.where(classes == y_i).item() else: classes_ext = torch.cat((classes._DNDarray__array, y_i._DNDarray__array.unsqueeze(0))) i = torch.argsort(classes_ext)[-1].item() where_y_i = ht.where(y == y_i)._DNDarray__array.tolist() X_i = X[where_y_i, :] if sample_weight is not None: sw_i = sample_weight[where_y_i] if 0 not in sw_i.shape: N_i = sw_i.sum() else: N_i = 0.0 sw_i = None else: sw_i = None N_i = X_i.shape[0] new_theta, new_sigma = self.__update_mean_variance( self.class_count_[i], self.theta_[i, :], self.sigma_[i, :], X_i, sw_i) self.theta_[i, :] = new_theta self.sigma_[i, :] = new_sigma self.class_count_[i] += N_i self.sigma_[:, :] += self.epsilon_ # Update if only no priors is provided if self.priors is None: # Empirical prior, with sample_weight taken into account self.class_prior_ = self.class_count_ / self.class_count_.sum() return self
def test_data_parallel(self): import heat.nn.functional as F with self.assertRaises(TypeError): ht.utils.data.datatools.DataLoader("asdf") class Model(ht.nn.Module): def __init__(self): super(Model, self).__init__() # 1 input image channel, 6 output channels, 3x3 square convolution # kernel self.conv1 = ht.nn.Conv2d(1, 6, 3) self.conv2 = ht.nn.Conv2d(6, 16, 3) # an affine operation: y = Wx + b self.fc1 = ht.nn.Linear(16 * 6 * 6, 120) # 6*6 from image dimension self.fc2 = ht.nn.Linear(120, 84) self.fc3 = ht.nn.Linear(84, 10) def forward(self, x): # Max pooling over a (2, 2) window x = self.conv1(x) x = F.max_pool2d(F.relu(x), (2, 2)) # If the size is a square you can only specify a single number x = F.max_pool2d(F.relu(self.conv2(x)), 2) x = x.view(-1, self.num_flat_features(x)) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x def num_flat_features(self, x): size = x.size()[1:] # all dimensions except the batch dimension num_features = 1 for s in size: num_features *= s return num_features class TestDataset(ht.utils.data.Dataset): def __init__(self, array, ishuffle): super(TestDataset, self).__init__(array, ishuffle=ishuffle) def __getitem__(self, item): return self.data[item] def Ishuffle(self): if not self.test_set: ht.utils.data.dataset_ishuffle(self, attrs=[["data", None]]) def Shuffle(self): if not self.test_set: ht.utils.data.dataset_shuffle(self, attrs=[["data", None]]) # create model and move it to GPU with id rank model = Model() optimizer = ht.optim.SGD(model.parameters(), lr=0.001) with self.assertRaises(TypeError): ht.optim.DataParallelOptimizer(optimizer, "asdf") dp_optimizer = ht.optim.DataParallelOptimizer(optimizer, True) ht.random.seed(1) torch.random.manual_seed(1) labels = torch.randn((2, 10), device=ht.get_device().torch_device) data = ht.random.rand(2 * ht.MPI_WORLD.size, 1, 32, 32, split=0) dataset = TestDataset(data, ishuffle=True) dataloader = ht.utils.data.datatools.DataLoader(dataset=dataset, batch_size=2) # there is only 1 batch on each process (data size[0] is 2 * number of processes, and the batch size is 2) self.assertTrue(len(dataloader) == 1) ht_model = ht.nn.DataParallel( model, data.comm, dp_optimizer, blocking_parameter_updates=True ) if str(ht.get_device())[:3] == "gpu": ht_model.to(ht.get_device().torch_device) lim = 1e-4 loss_fn = torch.nn.MSELoss() for _ in range(2): for data in dataloader: self.assertEqual(data.shape[0], 2) dp_optimizer.zero_grad() ht_outputs = ht_model(data) loss_fn(ht_outputs, labels).backward() dp_optimizer.step() for p in ht_model.parameters(): p0dim = p.shape[0] hld = ht.resplit(ht.array(p, is_split=0))._DNDarray__array hld_list = [hld[i * p0dim : (i + 1) * p0dim] for i in range(ht.MPI_WORLD.size - 1)] for i in range(1, len(hld_list)): self.assertTrue(torch.allclose(hld_list[0], hld_list[i], rtol=lim, atol=lim)) model = Model() optimizer = ht.optim.SGD(model.parameters(), lr=0.001) dp_optimizer = ht.optim.DataParallelOptimizer(optimizer, False) labels = torch.randn((2, 10), device=ht.get_device().torch_device) data = ht.random.rand(2 * ht.MPI_WORLD.size, 1, 32, 32, split=0) dataset = ht.utils.data.Dataset(data, ishuffle=False) dataloader = ht.utils.data.datatools.DataLoader(dataset=dataset, batch_size=2) ht_model = ht.nn.DataParallel( model, data.comm, dp_optimizer, blocking_parameter_updates=False ) if str(ht.get_device())[:3] == "gpu": ht_model.to(ht.get_device().torch_device) with self.assertRaises(TypeError): ht.nn.DataParallel(model, data.comm, "asdf") loss_fn = torch.nn.MSELoss() for _ in range(2): for data in dataloader: self.assertEqual(data.shape[0], 2) dp_optimizer.zero_grad() ht_outputs = ht_model(data) loss_fn(ht_outputs, labels).backward() dp_optimizer.step() for p in ht_model.parameters(): p0dim = p.shape[0] hld = ht.resplit(ht.array(p, is_split=0))._DNDarray__array hld_list = [hld[i * p0dim : (i + 1) * p0dim] for i in range(ht.MPI_WORLD.size - 1)] for i in range(1, len(hld_list)): self.assertTrue(torch.allclose(hld_list[0], hld_list[i], rtol=lim, atol=lim)) model = Model() optimizer = ht.optim.SGD(model.parameters(), lr=0.001) dp_optimizer = ht.optim.DataParallelOptimizer(optimizer, False) labels = torch.randn((2, 10), device=ht.get_device().torch_device) data = ht.random.rand(2 * ht.MPI_WORLD.size, 1, 32, 32, split=0) dataset = ht.utils.data.Dataset(data, ishuffle=True) dataloader = ht.utils.data.datatools.DataLoader(dataset=dataset, batch_size=2) ht_model = ht.nn.DataParallel( model, data.comm, dp_optimizer, blocking_parameter_updates=False ) if str(ht.get_device())[:3] == "gpu": ht_model.to(ht.get_device().torch_device) for _ in range(2): for data in dataloader: self.assertEqual(data.shape[0], 2) dp_optimizer.zero_grad() ht_outputs = ht_model(data) loss_fn(ht_outputs, labels).backward() dp_optimizer.step() for p in ht_model.parameters(): p0dim = p.shape[0] hld = ht.resplit(ht.array(p, is_split=0))._DNDarray__array hld_list = [hld[i * p0dim : (i + 1) * p0dim] for i in range(ht.MPI_WORLD.size - 1)] for i in range(1, len(hld_list)): self.assertTrue(torch.allclose(hld_list[0], hld_list[i], rtol=lim, atol=lim)) with self.assertWarns(Warning): ht_model = ht.nn.DataParallel( model, ht.MPI_WORLD, [dp_optimizer, dp_optimizer], blocking_parameter_updates=False ) # NOTE: this will throw a warning: this is expected self.assertTrue(ht_model.blocking_parameter_updates)