idx = np.random.permutation(n)
idx_train = idx[:int(0.6 * n)]
idx_val = idx[int(0.6 * n):int(0.8 * n)]
idx_test = idx[int(0.8 * n):]
# Transform the numpy matrices/vectors to torch tensors
features = torch.FloatTensor(features)
y = torch.LongTensor(np.argmax(class_labels, axis=1))
adj = torch.FloatTensor(adj)
idx_train = torch.LongTensor(idx_train)
idx_val = torch.LongTensor(idx_val)
idx_test = torch.LongTensor(idx_test)
# Creates the model and specifies the optimizer
model = GNN(features.shape[1], n_hidden_1, n_hidden_2, n_class, dropout_rate)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
def train(epoch):
t = time.time()
model.train()
optimizer.zero_grad()
output, _ = model(features, adj)
loss_train = F.nll_loss(output[idx_train], y[idx_train])
acc_train = accuracy(output[idx_train], y[idx_train])
loss_train.backward()
optimizer.step()
model.eval()
output, _ = model(features, adj)
class Trainer:
def __init__(self, params):
self.params = params
self.prj_path = Path(__file__).parent.resolve()
self.save_path = self.prj_path / 'pretrained' / f'{self.params.species}' / 'models'
if not self.save_path.exists():
self.save_path.mkdir(parents=True)
self.device = torch.device('cpu' if self.params.gpu ==
-1 else f'cuda:{params.gpu}')
self.num_cells, self.num_genes, self.num_labels, self.graph, self.train_ids, self.test_ids, self.labels = load_data_internal(
params)
self.labels = self.labels.to(self.device)
self.model = GNN(in_feats=self.params.dense_dim,
n_hidden=self.params.hidden_dim,
n_classes=self.num_labels,
n_layers=self.params.n_layers,
gene_num=self.num_genes,
activation=F.relu,
dropout=self.params.dropout).to(self.device)
self.optimizer = torch.optim.Adam(
self.model.parameters(),
lr=self.params.lr,
weight_decay=self.params.weight_decay)
self.loss_fn = nn.CrossEntropyLoss(reduction='sum')
if self.params.num_neighbors == 0:
self.num_neighbors = self.num_cells + self.num_genes
else:
self.num_neighbors = self.params.num_neighbors
print(
f"Train Number: {len(self.train_ids)}, Test Number: {len(self.test_ids)}"
)
def fit(self):
max_test_acc, _train_acc, _epoch = 0, 0, 0
for epoch in range(self.params.n_epochs):
loss = self.train()
train_correct, train_unsure = self.evaluate(
self.train_ids, 'train')
train_acc = train_correct / len(self.train_ids)
test_correct, test_unsure = self.evaluate(self.test_ids, 'test')
test_acc = test_correct / len(self.test_ids)
if max_test_acc <= test_acc:
final_test_correct_num = test_correct
final_test_unsure_num = test_unsure
_train_acc = train_acc
_epoch = epoch
max_test_acc = test_acc
self.save_model()
print(
f">>>>Epoch {epoch:04d}: Train Acc {train_acc:.4f}, Loss {loss / len(self.train_ids):.4f}, Test correct {test_correct}, "
f"Test unsure {test_unsure}, Test Acc {test_acc:.4f}")
if train_acc == 1:
break
print(
f"---{self.params.species} {self.params.tissue} Best test result:---"
)
print(
f"Epoch {_epoch:04d}, Train Acc {_train_acc:.4f}, Test Correct Num {final_test_correct_num}, Test Total Num {len(self.test_ids)}, Test Unsure Num {final_test_unsure_num}, Test Acc {final_test_correct_num / len(self.test_ids):.4f}"
)
def train(self):
self.model.train()
total_loss = 0
for batch, nf in enumerate(
NeighborSampler(g=self.graph,
batch_size=self.params.batch_size,
expand_factor=self.num_neighbors,
num_hops=self.params.n_layers,
neighbor_type='in',
shuffle=True,
num_workers=8,
seed_nodes=self.train_ids)):
nf.copy_from_parent(
) # Copy node/edge features from the parent graph.
logits = self.model(nf)
batch_nids = nf.layer_parent_nid(-1).type(
torch.long).to(device=self.device)
loss = self.loss_fn(logits, self.labels[batch_nids])
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
total_loss += loss.item()
return total_loss
def evaluate(self, ids, type='test'):
self.model.eval()
total_correct, total_unsure = 0, 0
for nf in NeighborSampler(g=self.graph,
batch_size=self.params.batch_size,
expand_factor=self.num_cells +
self.num_genes,
num_hops=params.n_layers,
neighbor_type='in',
shuffle=True,
num_workers=8,
seed_nodes=ids):
nf.copy_from_parent(
) # Copy node/edge features from the parent graph.
with torch.no_grad():
logits = self.model(nf).cpu()
batch_nids = nf.layer_parent_nid(-1).type(torch.long)
logits = nn.functional.softmax(logits, dim=1).numpy()
label_list = self.labels.cpu()[batch_nids]
for pred, label in zip(logits, label_list):
max_prob = pred.max().item()
if max_prob < self.params.unsure_rate / self.num_labels:
total_unsure += 1
elif pred.argmax().item() == label:
total_correct += 1
return total_correct, total_unsure
def save_model(self):
state = {
'model': self.model.state_dict(),
'optimizer': self.optimizer.state_dict()
}
torch.save(
state,
self.save_path / f"{self.params.species}-{self.params.tissue}.pt")
def main(_run, _config, _log):
'''
_config: dictionary; its keys and values are the variables setting in the cfg function
_run: run object defined by Sacred, can be used to record hashable values and get some information, e.g. run id, for a run
_log: logger object provided by Sacred, but is not very flexible, we can define loggers by oureselves
'''
config = dcopy(
_config
) # We need this step because Sacred does not allow us to change _config object
# But sometimes we need to add some key-value pairs to config
torch.cuda.set_device(config['gpu_id'])
save_source(_run) # Source code are saved by running this line
init_seed(config['seed'])
logger = init_logger(log_root=_run.observers[0].dir, file_name='log.txt')
output_folder_path = opjoin(_run.observers[0].dir,
config['path']['output_folder_name'])
os.makedirs(output_folder_path, exist_ok=True)
best_acc_list = []
last_acc_list = []
train_best_list = []
train_last_list = []
best_epoch = []
data = load_data(config=config)
split_iterator = range(config['data']['random_split']['num_splits']) \
if config['data']['random_split']['use'] \
else range(1)
config['adj'] = data[0]
for i in split_iterator:
output_folder = opjoin(output_folder_path, str(i))
os.makedirs(output_folder, exist_ok=True)
if config['data']['random_split']['use']:
data = resplit(
dataset=config['data']['dataset'],
data=data,
full_sup=config['data']['full_sup'],
num_classes=torch.unique(data[2]).shape[0],
num_nodes=data[1].shape[0],
num_per_class=config['data']['label_per_class'],
)
print(torch.sum(data[3]))
model = GNN(config=config)
if i == 0:
logger.info(model)
if config['use_gpu']:
model.cuda()
data = [
each.cuda() if hasattr(each, 'cuda') else each for each in data
]
optimizer = init_optimizer(
params=model.parameters(),
optim_type=config['optim']['type'],
lr=config['optim']['lr'],
weight_decay=config['optim']['weight_decay'],
momentum=config['optim']['momemtum'])
criterion = nn.NLLLoss()
best_model_path = opjoin(output_folder, 'best_model.pth')
last_model_path = opjoin(output_folder, 'last_model.pth')
best_dict_path = opjoin(output_folder, 'best_pred_dict.pkl')
last_dict_path = opjoin(output_folder, 'last_pred_dict.pkl')
losses_curve_path = opjoin(output_folder, 'losses.pkl')
accs_curve_path = opjoin(output_folder, 'accs.pkl')
best_state_path = opjoin(output_folder, 'best_state.pkl')
grads_path = opjoin(output_folder, 'grads.pkl')
best_pred_dict, last_pred_dict, train_losses, train_accs, \
val_losses, val_accs, best_state, grads, model_state = train(best_model_path,
last_model_path,
config,
criterion,
data,
logger,
model,
optimizer
)
last_model_state, best_model_state = model_state
losses_dict = {'train': train_losses, 'val': val_losses}
accs_dict = {'train': train_accs, 'val': val_accs}
logger.info(f'split_seed: {i: 04d}')
logger.info(f'Test set results on the last model:')
last_pred_dict = test(
criterion,
data,
last_model_path,
last_pred_dict,
logger,
model,
last_model_state,
)
logger.info(f'Test set results on the best model:')
if config['fastmode']:
best_pred_dict = last_pred_dict
else:
best_pred_dict = test(
criterion,
data,
best_model_path,
best_pred_dict,
logger,
model,
best_model_state,
)
logger.info('\n')
check_before_pkl(best_pred_dict)
with open(best_dict_path, 'wb') as f:
pkl.dump(best_pred_dict, f)
check_before_pkl(last_pred_dict)
with open(last_dict_path, 'wb') as f:
pkl.dump(last_pred_dict, f)
check_before_pkl(losses_dict)
with open(losses_curve_path, 'wb') as f:
pkl.dump(losses_dict, f)
check_before_pkl(accs_dict)
with open(accs_curve_path, 'wb') as f:
pkl.dump(accs_dict, f)
check_before_pkl(best_state)
with open(best_state_path, 'wb') as f:
pkl.dump(best_state, f)
check_before_pkl(grads)
with open(grads_path, 'wb') as f:
pkl.dump(grads, f)
best_acc_list.append(best_pred_dict['test acc'].item())
last_acc_list.append(last_pred_dict['test acc'].item())
train_best_list.append(best_state['train acc'].item())
train_last_list.append(train_accs[-1].item())
best_epoch.append(best_state['epoch'])
logger.info('********************* STATISTICS *********************')
np.set_printoptions(precision=4, suppress=True)
logger.info(f"\n"
f"Best test acc: {best_acc_list}\n"
f"Mean: {np.mean(best_acc_list)}\t"
f"Std: {np.std(best_acc_list)}\n"
f"Last test acc: {last_acc_list}\n"
f"Mean: {np.mean(last_acc_list)}\t"
f"Std: {np.std(last_acc_list)}\n")
logger.info(f"best epoch: {best_epoch}")
