AlphaFold3 data_modules 模塊的 OpenFoldDataset
類的 looped_samples 方法用于 循環采樣數據,確保數據能被不斷地提供,適用于 PyTorch 的 DataLoader
在訓練過程中迭代讀取數據。dataset_idx
指定了當前要處理的數據集(即 self.datasets[dataset_idx]
)
源代碼:
def looped_samples(self, dataset_idx):max_cache_len = int(self.epoch_len * self.probabilities[dataset_idx])dataset = self.datasets[dataset_idx]idx_iter = self.looped_shuffled_dataset_idx(len(dataset))chain_data_cache = dataset.chain_data_cachewhile True:weights = []idx = []for _ in range(max_cache_len):candidate_idx = next(idx_iter)chain_id = dataset.idx_to_chain_id(candidate_idx)chain_data_cache_entry = chain_data_cache[chain_id]if not self.deterministic_train_filter(chain_data_cache_entry):continuep = self.get_stochastic_train_filter_prob(chain_data_cache_entry,)weights.append([1. - p, p])idx.append(candidate_idx)samples = torch.multinomial(torch.tensor(weights),num_samples=1,generator=self.generator,)samples = samples.squeeze()cache = [i for i, s in zip(idx, samples) if s]for datapoint_idx in cache:yield datapoint_idx
源碼解讀:
max_cache_len = int(self.epoch_len * self.probabilities[dataset_idx])
-
epoch_len
是一個訓練周期(epoch)中期望的樣本總數。 -
self.probabilitie