def main(args):
batch_size = args.batch_size
model_fname = args.mod_name
if multi_gpu and batch_size < torch.cuda.device_count():
exit('Batch size too small')
# make a folder for the graphs of this model
Path(args.output_dir).mkdir(exist_ok=True)
save_dir = osp.join(args.output_dir, model_fname)
Path(save_dir).mkdir(exist_ok=True)
# get dataset and split
gdata = GraphDataset(root=args.input_dir, bb=args.box_num)
# merge data from separate files into one contiguous array
bag = []
for g in gdata:
bag += g
random.Random(0).shuffle(bag)
bag = bag[:args.num_data]
# temporary patch to use px, py, pz
for d in bag:
d.x = d.x[:, :3]
# 80:10:10 split datasets
fulllen = len(bag)
train_len = int(0.8 * fulllen)
tv_len = int(0.10 * fulllen)
train_dataset = bag[:train_len]
valid_dataset = bag[train_len:train_len + tv_len]
test_dataset = bag[train_len + tv_len:]
train_samples = len(train_dataset)
valid_samples = len(valid_dataset)
test_samples = len(test_dataset)
if multi_gpu:
train_loader = DataListLoader(train_dataset,
batch_size=batch_size,
pin_memory=True,
shuffle=True)
valid_loader = DataListLoader(valid_dataset,
batch_size=batch_size,
pin_memory=True,
shuffle=False)
test_loader = DataListLoader(test_dataset,
batch_size=batch_size,
pin_memory=True,
shuffle=False)
else:
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
pin_memory=True,
shuffle=True)
valid_loader = DataLoader(valid_dataset,
batch_size=batch_size,
pin_memory=True,
shuffle=False)
test_loader = DataLoader(test_dataset,
batch_size=batch_size,
pin_memory=True,
shuffle=False)
# specify loss function
loss_ftn_obj = LossFunction(args.loss,
emd_modname=args.emd_model_name,
device=device)
# create model
input_dim = 3
big_dim = 32
hidden_dim = args.lat_dim
lr = args.lr
patience = args.patience
if args.model == 'MetaLayerGAE':
model = models.GNNAutoEncoder()
else:
if args.model[-3:] == 'EMD':
model = getattr(models,
args.model)(input_dim=input_dim,
big_dim=big_dim,
hidden_dim=hidden_dim,
emd_modname=args.emd_model_name)
else:
model = getattr(models, args.model)(input_dim=input_dim,
big_dim=big_dim,
hidden_dim=hidden_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=4)
valid_losses = []
train_losses = []
start_epoch = 0
n_epochs = 200
# load in model
modpath = osp.join(save_dir, model_fname + '.best.pth')
try:
model.load_state_dict(torch.load(modpath))
train_losses, valid_losses, start_epoch = torch.load(
osp.join(save_dir, 'losses.pt'))
print('Loaded model')
best_valid_loss = test(model, valid_loader, valid_samples, batch_size,
loss_ftn_obj)
print(f'Saved model valid loss: {best_valid_loss}')
except:
print('Creating new model')
best_valid_loss = 9999999
if multi_gpu:
model = DataParallel(model)
model.to(torch.device(device))
# Training loop
stale_epochs = 0
loss = best_valid_loss
for epoch in range(start_epoch, n_epochs):
if multi_gpu:
loss = train_parallel(model, optimizer, train_loader,
train_samples, batch_size, loss_ftn_obj)
valid_loss = test_parallel(model, valid_loader, valid_samples,
batch_size, loss_ftn_obj)
else:
loss = train(model, optimizer, train_loader, train_samples,
batch_size, loss_ftn_obj)
valid_loss = test(model, valid_loader, valid_samples, batch_size,
loss_ftn_obj)
scheduler.step(valid_loss)
train_losses.append(loss)
valid_losses.append(valid_loss)
print('Epoch: {:02d}, Training Loss: {:.4f}'.format(epoch, loss))
print(' Validation Loss: {:.4f}'.format(valid_loss))
if valid_loss < best_valid_loss:
best_valid_loss = valid_loss
print('New best model saved to:', modpath)
if multi_gpu:
torch.save(model.module.state_dict(), modpath)
else:
torch.save(model.state_dict(), modpath)
torch.save((train_losses, valid_losses, epoch + 1),
osp.join(save_dir, 'losses.pt'))
stale_epochs = 0
else:
stale_epochs += 1
print(
f'Stale epoch: {stale_epochs}\nBest: {best_valid_loss}\nCurr: {valid_loss}'
)
if stale_epochs >= patience:
print('Early stopping after %i stale epochs' % patience)
break
# model training done
train_epochs = list(range(epoch + 1))
early_stop_epoch = epoch - stale_epochs
loss_curves(train_epochs, early_stop_epoch, train_losses, valid_losses,
save_dir)
# compare input and reconstructions
model.load_state_dict(torch.load(modpath))
input_fts = []
reco_fts = []
for t in valid_loader:
model.eval()
if isinstance(t, list):
for d in t:
input_fts.append(d.x)
else:
input_fts.append(t.x)
t.to(device)
reco_out = model(t)
if isinstance(reco_out, tuple):
reco_out = reco_out[0]
reco_fts.append(reco_out.cpu().detach())
input_fts = torch.cat(input_fts)
reco_fts = torch.cat(reco_fts)
plot_reco_difference(
input_fts, reco_fts, model_fname,
osp.join(save_dir, 'reconstruction_post_train', 'valid'))
input_fts = []
reco_fts = []
for t in test_loader:
model.eval()
if isinstance(t, list):
for d in t:
input_fts.append(d.x)
else:
input_fts.append(t.x)
t.to(device)
reco_out = model(t)
if isinstance(reco_out, tuple):
reco_out = reco_out[0]
reco_fts.append(reco_out.cpu().detach())
input_fts = torch.cat(input_fts)
reco_fts = torch.cat(reco_fts)
plot_reco_difference(
input_fts, reco_fts, model_fname,
osp.join(save_dir, 'reconstruction_post_train', 'test'))
print('Completed')
def train_cross_validation(model_cls,
dataset,
dropout=0.0,
lr=1e-3,
weight_decay=1e-2,
num_epochs=200,
n_splits=10,
use_gpu=True,
dp=False,
ddp=False,
comment='',
tb_service_loc='192.168.192.57:6007',
batch_size=1,
num_workers=0,
pin_memory=False,
cuda_device=None,
tb_dir='runs',
model_save_dir='saved_models',
res_save_dir='res',
fold_no=None,
saved_model_path=None,
device_ids=None,
patience=20,
seed=None,
fold_seed=None,
save_model=False,
is_reg=True,
live_loss=True,
domain_cls=True,
final_cls=True):
"""
:type fold_seed: int
:param live_loss: bool
:param is_reg: bool
:param save_model: bool
:param seed:
:param patience: for early stopping
:param device_ids: for ddp
:param saved_model_path:
:param fold_no: int
:param ddp_port: str
:param ddp: DDP
:param cuda_device: list of int
:param pin_memory: bool, DataLoader args
:param num_workers: int, DataLoader args
:param model_cls: pytorch Module cls
:param dataset: instance
:param dropout: float
:param lr: float
:param weight_decay:
:param num_epochs:
:param n_splits: number of kFolds
:param use_gpu: bool
:param dp: bool
:param comment: comment in the logs, to filter runs in tensorboard
:param tb_service_loc: tensorboard service location
:param batch_size: Dataset args not DataLoader
:return:
"""
saved_args = locals()
seed = int(time.time() % 1e4 * 1e5) if seed is None else seed
saved_args['random_seed'] = seed
torch.manual_seed(seed)
np.random.seed(seed)
if use_gpu:
torch.cuda.manual_seed_all(seed)
# torch.backends.cudnn.deterministic = True
# torch.backends.cudnn.benchmark = False
model_name = model_cls.__name__
if not cuda_device:
if device_ids and dp:
device = device_ids[0]
else:
device = torch.device(
'cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
else:
device = cuda_device
device_count = torch.cuda.device_count() if dp else 1
device_count = len(device_ids) if (device_ids is not None
and dp) else device_count
batch_size = batch_size * device_count
# TensorBoard
log_dir_base = get_model_log_dir(comment, model_name)
if tb_service_loc is not None:
print("TensorBoard available at http://{1}/#scalars®exInput={0}".
format(log_dir_base, tb_service_loc))
else:
print("Please set up TensorBoard")
# model
criterion = nn.NLLLoss()
print("Training {0} {1} models for cross validation...".format(
n_splits, model_name))
# 1
# folds, fold = KFold(n_splits=n_splits, shuffle=False, random_state=seed), 0
# 2
# folds = GroupKFold(n_splits=n_splits)
# iter = folds.split(np.zeros(len(dataset)), groups=dataset.data.site_id)
# 4
# folds = StratifiedKFold(n_splits=n_splits, random_state=fold_seed, shuffle=True if fold_seed else False)
# iter = folds.split(np.zeros(len(dataset)), dataset.data.y.numpy(), groups=dataset.data.subject_id)
# 5
fold = 0
iter = multi_site_cv_split(dataset.data.y,
dataset.data.site_id,
dataset.data.subject_id,
n_splits,
random_state=fold_seed,
shuffle=True if fold_seed else False)
for train_idx, val_idx in tqdm_notebook(iter, desc='CV', leave=False):
fold += 1
liveloss = PlotLosses() if live_loss else None
# for a specific fold
if fold_no is not None:
if fold != fold_no:
continue
writer = SummaryWriter(log_dir=osp.join('runs', log_dir_base +
str(fold)))
model_save_dir = osp.join('saved_models', log_dir_base + str(fold))
print("creating dataloader tor fold {}".format(fold))
train_dataset, val_dataset = norm_train_val(dataset, train_idx,
val_idx)
model = model_cls(writer)
train_dataloader = DataLoader(train_dataset,
shuffle=True,
batch_size=batch_size,
collate_fn=lambda data_list: data_list,
num_workers=num_workers,
pin_memory=pin_memory)
val_dataloader = DataLoader(val_dataset,
shuffle=False,
batch_size=batch_size,
collate_fn=lambda data_list: data_list,
num_workers=num_workers,
pin_memory=pin_memory)
if fold == 1 or fold_no is not None:
print(model)
writer.add_text('model_summary', model.__repr__())
writer.add_text('training_args', str(saved_args))
optimizer = torch.optim.AdamW(model.parameters(),
lr=lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=weight_decay,
amsgrad=False)
# scheduler_reduce = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5)
scheduler = GradualWarmupScheduler(optimizer,
multiplier=10,
total_epoch=5)
# scheduler = scheduler_reduce
# optimizer = torch.optim.SGD(model.parameters(), lr=lr, weight_decay=weight_decay)
if dp and use_gpu:
model = model.cuda() if device_ids is None else model.to(
device_ids[0])
model = DataParallel(model, device_ids=device_ids)
elif use_gpu:
model = model.to(device)
if saved_model_path is not None:
model.load_state_dict(torch.load(saved_model_path))
best_map, patience_counter, best_score = 0.0, 0, np.inf
for epoch in tqdm_notebook(range(1, num_epochs + 1),
desc='Epoch',
leave=False):
logs = {}
# scheduler.step(epoch=epoch, metrics=best_score)
for phase in ['train', 'validation']:
if phase == 'train':
model.train()
dataloader = train_dataloader
else:
model.eval()
dataloader = val_dataloader
# Logging
running_total_loss = 0.0
running_corrects = 0
running_reg_loss = 0.0
running_nll_loss = 0.0
epoch_yhat_0, epoch_yhat_1 = torch.tensor([]), torch.tensor([])
epoch_label, epoch_predicted = torch.tensor([]), torch.tensor(
[])
logging_hist = True if phase == 'train' else False # once per epoch
for data_list in tqdm_notebook(dataloader,
desc=phase,
leave=False):
# TODO: check devices
if dp:
data_list = to_cuda(data_list,
(device_ids[0] if device_ids
is not None else 'cuda'))
y_hat, domain_yhat, reg = model(data_list)
y = torch.tensor([],
dtype=dataset.data.y.dtype,
device=device)
domain_y = torch.tensor([],
dtype=dataset.data.site_id.dtype,
device=device)
for data in data_list:
y = torch.cat([y, data.y.view(-1).to(device)])
domain_y = torch.cat(
[domain_y,
data.site_id.view(-1).to(device)])
loss = criterion(y_hat, y)
domain_loss = criterion(domain_yhat, domain_y)
# domain_loss = -1e-7 * domain_loss
# print(domain_loss.item())
if domain_cls:
total_loss = domain_loss
_, predicted = torch.max(domain_yhat, 1)
label = domain_y
if final_cls:
total_loss = loss
_, predicted = torch.max(y_hat, 1)
label = y
if domain_cls and final_cls:
total_loss = (loss + domain_loss).sum()
_, predicted = torch.max(y_hat, 1)
label = y
if is_reg:
total_loss += reg.sum()
if phase == 'train':
# print(torch.autograd.grad(y_hat.sum(), model.saved_x, retain_graph=True))
optimizer.zero_grad()
total_loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 2.0)
optimizer.step()
running_nll_loss += loss.item()
running_total_loss += total_loss.item()
running_reg_loss += reg.sum().item()
running_corrects += (predicted == label).sum().item()
epoch_yhat_0 = torch.cat(
[epoch_yhat_0, y_hat[:, 0].detach().view(-1).cpu()])
epoch_yhat_1 = torch.cat(
[epoch_yhat_1, y_hat[:, 1].detach().view(-1).cpu()])
epoch_label = torch.cat(
[epoch_label,
label.detach().float().view(-1).cpu()])
epoch_predicted = torch.cat([
epoch_predicted,
predicted.detach().float().view(-1).cpu()
])
# precision = sklearn.metrics.precision_score(epoch_label, epoch_predicted, average='micro')
# recall = sklearn.metrics.recall_score(epoch_label, epoch_predicted, average='micro')
# f1_score = sklearn.metrics.f1_score(epoch_label, epoch_predicted, average='micro')
accuracy = sklearn.metrics.accuracy_score(
epoch_label, epoch_predicted)
epoch_total_loss = running_total_loss / dataloader.__len__()
epoch_nll_loss = running_nll_loss / dataloader.__len__()
epoch_reg_loss = running_reg_loss / dataloader.__len__()
# print('epoch {} {}_nll_loss: {}'.format(epoch, phase, epoch_nll_loss))
writer.add_scalars(
'nll_loss', {'{}_nll_loss'.format(phase): epoch_nll_loss},
epoch)
writer.add_scalars('accuracy',
{'{}_accuracy'.format(phase): accuracy},
epoch)
# writer.add_scalars('{}_APRF'.format(phase),
# {
# 'accuracy': accuracy,
# 'precision': precision,
# 'recall': recall,
# 'f1_score': f1_score
# },
# epoch)
if epoch_reg_loss != 0:
writer.add_scalars(
'reg_loss'.format(phase),
{'{}_reg_loss'.format(phase): epoch_reg_loss}, epoch)
# print(epoch_reg_loss)
# writer.add_histogram('hist/{}_yhat_0'.format(phase),
# epoch_yhat_0,
# epoch)
# writer.add_histogram('hist/{}_yhat_1'.format(phase),
# epoch_yhat_1,
# epoch)
# Save Model & Early Stopping
if phase == 'validation':
model_save_path = model_save_dir + '-{}-{}-{:.3f}-{:.3f}'.format(
model_name, epoch, accuracy, epoch_nll_loss)
# best score
if accuracy > best_map:
best_map = accuracy
model_save_path = model_save_path + '-best'
score = epoch_nll_loss
if score < best_score:
patience_counter = 0
best_score = score
else:
patience_counter += 1
# skip first 10 epoch
# best_score = best_score if epoch > 10 else -np.inf
if save_model:
for th, pfix in zip(
[0.8, 0.75, 0.7, 0.5, 0.0],
['-perfect', '-great', '-good', '-bad', '-miss']):
if accuracy >= th:
model_save_path += pfix
break
torch.save(model.state_dict(), model_save_path)
writer.add_scalars('best_val_accuracy',
{'{}_accuracy'.format(phase): best_map},
epoch)
writer.add_scalars(
'best_nll_loss',
{'{}_nll_loss'.format(phase): best_score}, epoch)
writer.add_scalars('learning_rate', {
'learning_rate':
scheduler.optimizer.param_groups[0]['lr']
}, epoch)
if patience_counter >= patience:
print("Stopped at epoch {}".format(epoch))
return
if live_loss:
prefix = ''
if phase == 'validation':
prefix = 'val_'
logs[prefix + 'log loss'] = epoch_nll_loss
logs[prefix + 'accuracy'] = accuracy
if live_loss:
liveloss.update(logs)
liveloss.draw()
print("Done !")
def test(args):
model_train_dict = torch.load(args.checkpoint)
model = GeometricDataParallel(
PotentialNetParallel(
in_channels=20,
out_channels=1,
covalent_gather_width=model_train_dict["args"]
["covalent_gather_width"],
non_covalent_gather_width=model_train_dict["args"]
["non_covalent_gather_width"],
covalent_k=model_train_dict["args"]["covalent_k"],
non_covalent_k=model_train_dict["args"]["non_covalent_k"],
covalent_neighbor_threshold=model_train_dict["args"]
["covalent_threshold"],
non_covalent_neighbor_threshold=model_train_dict["args"]
["non_covalent_threshold"],
)).float()
model.load_state_dict(model_train_dict["model_state_dict"])
dataset_list = []
# because the script allows for multiple datasets, we iterate over the list of files to build one combined dataset object
for data in args.test_data:
dataset_list.append(
PDBBindDataset(
data_file=data,
dataset_name=args.dataset_name,
feature_type=args.feature_type,
preprocessing_type=args.preprocessing_type,
output_info=True,
cache_data=False,
use_docking=args.use_docking,
))
dataset = ConcatDataset(dataset_list)
print("{} complexes in dataset".format(len(dataset)))
dataloader = DataListLoader(dataset,
batch_size=args.batch_size,
shuffle=False)
model.eval()
model.cuda()
if args.print_model:
print(model)
print("{} total parameters.".format(
sum(p.numel() for p in model.parameters())))
if not os.path.exists(args.output):
os.makedirs(args.output)
output_f = "{}/{}".format(args.output, args.output_file_name)
with h5py.File(output_f, "w") as f:
for batch in tqdm(dataloader):
batch = [x for x in batch if x is not None]
if len(batch) < 1:
continue
for item in batch:
name = item[0]
pose = item[1]
data = item[2]
name_grp = f.require_group(str(name))
name_pose_grp = name_grp.require_group(str(pose))
y = data.y
name_pose_grp.attrs["y_true"] = y
(
covalent_feature,
non_covalent_feature,
pool_feature,
fc0_feature,
fc1_feature,
y_,
) = model.module(Batch().from_data_list([data]),
return_hidden_feature=True)
name_pose_grp.attrs["y_pred"] = y_.cpu().data.numpy()
hidden_features = np.concatenate(
(
covalent_feature.cpu().data.numpy(),
non_covalent_feature.cpu().data.numpy(),
pool_feature.cpu().data.numpy(),
fc0_feature.cpu().data.numpy(),
fc1_feature.cpu().data.numpy(),
),
axis=1,
)
name_pose_grp.create_dataset(
"hidden_features",
(hidden_features.shape[0], hidden_features.shape[1]),
data=hidden_features,
)
class Trainer():
def __init__(self,
option,
model,
train_dataset,
valid_dataset,
test_dataset=None):
self.option = option
self.device = torch.device("cuda:{}".format(option['cuda_devices'][0]) \
if torch.cuda.is_available() else "cpu")
self.model = DataParallel(model, device_ids=self.option['cuda_devices']).to(self.device) \
if option['parallel'] else model.to(self.device)
# Setting the train valid and test data loader
if self.option['parallel']:
self.train_dataloader = DataListLoader(train_dataset, \
batch_size=self.option['train_batch'])
self.valid_dataloader = DataListLoader(valid_dataset,
batch_size=64)
if test_dataset:
self.test_dataloader = DataListLoader(test_dataset,
batch_size=64)
else:
self.train_dataloader = DataLoader(train_dataset, \
batch_size=self.option['train_batch'])
self.valid_dataloader = DataLoader(valid_dataset, batch_size=64)
if test_dataset:
self.test_dataloader = DataLoader(test_dataset, batch_size=64)
# Setting the Adam optimizer with hyper-param
self.criterion = torch.nn.L1Loss()
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.option['lr'])
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer,
mode='min',
factor=0.7,
patience=self.option['lr_scheduler_patience'],
min_lr=0.0000001)
# other
self.start = time.time()
self.save_id = ''.join(
random.sample('zyxwvutsrqponmlkjihgfedcba1234567890', 4))
self.abs_file_dir = os.path.dirname(os.path.abspath(__file__))
self.ckpt_save_dir = os.path.join(
self.abs_file_dir, 'ckpt',
'ckpts_task{}_{}'.format(self.option['task'], self.save_id))
self.log_save_path = os.path.join(
self.abs_file_dir, 'log',
'log_task{}_{}.txt'.format(self.option['task'], self.save_id))
self.record_save_path = os.path.join(
self.abs_file_dir, 'record',
'record_task{}_{}.csv'.format(self.option['task'], self.save_id))
os.system('mkdir -p log record {}'.format(self.ckpt_save_dir))
self.records = {
'trn_record': [],
'val_record': [],
'val_losses': [],
'best_ckpt': None
}
self.log(
msgs=['\t{}:{}\n'.format(k, v) for k, v in self.option.items()])
self.log('save id: {}'.format(self.save_id))
self.log('train set num:{} valid set num:{} test set num: {}'.format(
len(train_dataset), len(valid_dataset), len(test_dataset)))
self.log("Total Parameters:" +
str(sum([p.nelement() for p in self.model.parameters()])))
def train_iterations(self, epoch):
self.model.train()
losses = []
for i, data in enumerate(self.train_dataloader):
self.optimizer.zero_grad()
if self.option['parallel']:
sample_list = data # data will be sample_list in parallel model
output = self.model(sample_list)
y = torch.cat([sample.y
for sample in sample_list]).to(output.device)
loss = self.criterion(output, y)
else:
data = data.to(self.device)
output = self.model(data)
loss = self.criterion(output, data.y)
loss.backward()
self.optimizer.step()
losses.append(loss.item())
if i % 100 == 0:
self.log('\tbatch {} training loss: {:.5f}'.format(
i, loss.item()),
with_time=True)
trn_loss = np.array(losses).mean()
return trn_loss
def valid_iterations(self, epoch=None, mode='valid'):
self.model.eval()
if mode == 'test': dataloader = self.test_dataloader
if mode == 'valid': dataloader = self.valid_dataloader
outputs = []
ys = []
with torch.no_grad():
for data in dataloader:
if self.option['parallel']:
sample_list = data # data will be samplelist in parallel model
output = self.model(sample_list)
y = torch.cat([sample.y for sample in sample_list
]).to(output.device)
else:
data = data.to(self.device)
output = self.model(data)
y = data.y
outputs.append(output)
ys.append(y)
val_loss = self.criterion(torch.cat(outputs), torch.cat(ys)).item()
if mode == 'test': self.log('Test loss: {:.5f}'.format(val_loss))
return val_loss
def train(self):
self.log('Training start...')
early_stop_cnt = 0
for epoch in tqdm(range(self.option['train_epoch'])):
trn_loss = self.train_iterations(epoch)
val_loss = self.valid_iterations(epoch)
self.scheduler.step(val_loss)
lr_cur = self.scheduler.optimizer.param_groups[0]['lr']
self.log('Epoch:{} trn_loss:{:.5f} val_loss:{:.5f} lr_cur:{:.7f}'.
format(epoch, trn_loss, val_loss, lr_cur),
with_time=True)
self.records['val_losses'].append(val_loss)
self.records['val_record'].append([epoch, val_loss, lr_cur])
self.records['trn_record'].append([epoch, trn_loss, lr_cur])
if val_loss == np.array(self.records['val_losses']).min():
self.save_model_and_records(epoch, trn_loss, val_loss)
early_stop_cnt = 0
else:
early_stop_cnt += 1
if self.option[
'early_stop_patience'] > 0 and early_stop_cnt > self.option[
'early_stop_patience']:
self.log('Early stop hitted!')
break
self.save_model_and_records(epoch, trn_loss, val_loss, final_save=True)
def save_model_and_records(self,
epoch,
trn_loss,
val_loss,
final_save=False):
if final_save:
self.save_loss_records()
file_name = 'Final_save{}_{}_{:.5f}_{:.5f}.ckpt'.format(
self.option['task'], epoch, trn_loss, val_loss)
else:
file_name = 'task{}_{}_{:.5f}_{:.5f}.ckpt'.format(
self.option['task'], epoch, trn_loss, val_loss)
self.records['best_ckpt'] = file_name
with open(os.path.join(self.ckpt_save_dir, file_name), 'wb') as f:
torch.save(
{
'option': self.option,
'records': self.records,
'model_state_dict': self.model.state_dict(),
}, f)
self.log('Model saved at epoch {}'.format(epoch))
def save_loss_records(self):
trn_record = pd.DataFrame(self.records['trn_record'],
columns=['epoch', 'trn_loss', 'lr'])
val_record = pd.DataFrame(self.records['val_record'],
columns=['epoch', 'val_loss', 'lr'])
ret = pd.DataFrame({
'Epoch': trn_record['epoch'],
'Traning MAE Loss': trn_record['trn_loss'],
'Validation MAE Loss': val_record['val_loss'],
})
ret.to_csv(self.record_save_path)
return ret
def load_best_ckpt(self):
ckpt_path = self.ckpt_save_dir + '/' + self.records['best_ckpt']
self.log('The best ckpt is {}'.format(ckpt_path))
self.load_ckpt(ckpt_path)
def load_ckpt(self, ckpt_path):
self.log('Ckpt loading: {}'.format(ckpt_path))
ckpt = torch.load(ckpt_path)
self.option = ckpt['option']
self.records = ckpt['records']
self.model.load_state_dict(ckpt['model_state_dict'])
def log(self, msg=None, msgs=None, with_time=False):
if with_time:
msg = msg + ' Time elapsed {:.2f} hrs ({:.1f} mins)'.format(
(time.time() - self.start) / 3600.,
(time.time() - self.start) / 60.)
with open(self.log_save_path, 'a+') as f:
if msgs:
f.writelines(msgs)
for x in msgs:
print(x, end='')
if msg:
f.write(msg + '\n')
print(msg)
loss.append(tr_loss.detach().item())
tr_loss.backward()
optimizer.step()
if batch_n == 0:
tr_weights = weights
first_batch_labels = labels.clone().detach().to(device)
pred = out.clone().detach().round().to(device)
loss = mean(loss)
print("---- Round {}: loss={:.4f} lr:{:.6f}"
.format(epoch, loss, learn_rate))
# -------------- REPORTING ------------------------------------
roc_auc = roc_auc_score(first_batch_labels.cpu(), pred.cpu())
model.eval()
cum_pred = torch.Tensor().to(device)
cum_labels = torch.Tensor().to(device)
te_weights = torch.Tensor().to(device)
for batch_n, datalist in enumerate(val_loader):
out = model(datalist)
labels = torch.cat([data.y for data in datalist]).to(out.device)
weights = generate_weights(labels).to(out.device)
te_loss = F.binary_cross_entropy(out, target=labels, weight=generate_weights(labels))
pred = out.detach().round().to(device)
cum_labels = torch.cat((cum_labels, labels.clone().detach()), dim=0)
cum_pred = torch.cat((cum_pred, pred.clone().detach()), dim=0)
te_weights = torch.cat((te_weights, weights.clone().detach()), dim=0)
roc_auc_te = roc_auc_score(cum_labels.cpu(), cum_pred.cpu())
writer.add_scalars('Loss', {'train': tr_loss,
class Trainer():
def __init__(self,
option,
model,
train_dataset=None,
valid_dataset=None,
test_dataset=None,
weight=[[1.0, 1.0]],
tasks_num=1):
self.option = option
# self.tasks = ["MUV-466","MUV-548","MUV-600","MUV-644","MUV-652","MUV-689","MUV-692","MUV-712","MUV-713",
# "MUV-733","MUV-737","MUV-810","MUV-832","MUV-846","MUV-852","MUV-858","MUV-859"]
self.tasks_num = tasks_num
self.save_path = self.option['exp_path']
self.device = torch.device("cuda:{}".format(0) \
if torch.cuda.is_available() and not option['cpu'] else "cpu")
self.model = DataParallel(model).to(self.device) \
if option['parallel'] else model.to(self.device)
#Setting the train valid and test data loader
if train_dataset and valid_dataset:
if self.option['parallel']:
self.train_dataloader = DataListLoader(train_dataset, \
batch_size=self.option['batch_size'],shuffle=True)
self.valid_dataloader = DataListLoader(
valid_dataset, batch_size=self.option['batch_size'])
if test_dataset:
self.test_dataloader = DataListLoader(
test_dataset, batch_size=self.option['batch_size'])
else:
self.train_dataloader = DataLoader(train_dataset, \
batch_size=self.option['batch_size'],shuffle=True,num_workers=4)
self.valid_dataloader = DataLoader(
valid_dataset,
batch_size=self.option['batch_size'],
num_workers=4)
if test_dataset:
self.test_dataloader = DataLoader(
test_dataset,
batch_size=self.option['batch_size'],
num_workers=4)
else:
self.test_dataset = test_dataset
if self.option['parallel']:
self.test_dataloader = DataListLoader(
test_dataset,
batch_size=self.option['batch_size'],
num_workers=0)
else:
self.test_dataloader = DataLoader(
test_dataset,
batch_size=self.option['batch_size'],
num_workers=4)
# Setting the Adam optimizer with hyper-param
if not option['focalloss']:
self.criterion = [
torch.nn.CrossEntropyLoss(torch.Tensor(w).to(self.device),
reduction='mean') for w in weight
]
else:
self.log('Using FocalLoss')
self.criterion = [FocalLoss(alpha=1 / w[0])
for w in weight] #alpha 0.965
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.option['lr'],
weight_decay=option['weight_decay'])
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer,
mode='min',
factor=0.7,
patience=self.option['lr_scheduler_patience'],
min_lr=1e-6)
self.start = time.time()
self.records = {
'best_epoch': None,
'val_auc': [],
'best_val_auc': 0.,
'best_trn_auc': 0.,
'best_test_auc': 0.
}
self.log(
msgs=['\t{}:{}\n'.format(k, v) for k, v in self.option.items()],
show=False)
if train_dataset:
self.log(
'train set num:{} valid set num:{} test set num: {}'.
format(len(train_dataset), len(valid_dataset),
len(test_dataset)))
self.log("total parameters:" +
str(sum([p.nelement() for p in self.model.parameters()])))
self.log(msgs=str(model).split('\n'), show=False)
def train_iterations(self):
self.model.train()
losses = []
y_out_all = []
y_pred_list = {}
y_label_list = {}
for data in tqdm(self.train_dataloader):
self.optimizer.zero_grad()
if not self.option['parallel']:
data = data.to(self.device)
target_len = data.pr_len
y_idx = torch.zeros(target_len[0]).long()
for i, e in enumerate(target_len):
if i > 0:
y_idx = torch.cat(
[y_idx, torch.full((e.item(), ), i).long()])
y_idx = y_idx.to(self.device)
data.protein = torch_geometric.utils.to_dense_batch(
data.protein, y_idx)[0]
output = self.model(data)
else:
output = self.model(data)
data = Batch.from_data_list(data).to(self.device)
loss = 0
for i in range(self.tasks_num):
y_pred = output
y_label = data.y
loss += self.criterion[i](y_pred, y_label)
probs = F.softmax(y_pred.detach().cpu(), dim=-1)
y_out = probs.argmax(dim=1).numpy()
y_out_all.extend(y_out)
y_pred = probs[:, 1].view(-1).numpy()
# print(i,np.isnan(y_pred).any())
try:
y_label_list[i].extend(y_label.cpu().numpy())
y_pred_list[i].extend(y_pred)
except:
y_label_list[i] = []
y_pred_list[i] = []
y_label_list[i].extend(y_label.cpu().numpy())
y_pred_list[i].extend(y_pred)
# print(np.isnan(y_label_list[i]))
loss.backward()
self.optimizer.step()
losses.append(loss.item())
val_precision = metrics.precision_score(y_label_list[0],
np.array(y_out_all))
val_recall = metrics.recall_score(y_label_list[0], np.array(y_out_all))
trn_roc = [
metrics.roc_auc_score(y_label_list[i], y_pred_list[i])
for i in range(self.tasks_num)
]
trn_loss = np.array(losses).mean()
return trn_loss, np.array(trn_roc).mean(), val_precision, val_recall
def valid_iterations(self, mode='valid'):
self.model.eval()
if mode == 'test' or mode == 'eval': dataloader = self.test_dataloader
if mode == 'valid': dataloader = self.valid_dataloader
losses = []
y_out_all = []
y_pred_list = {}
y_label_list = {}
with torch.no_grad():
for data in tqdm(dataloader):
if not self.option['parallel']:
data = data.to(self.device)
target_len = data.pr_len
y_idx = torch.zeros(target_len[0]).long()
for i, e in enumerate(target_len):
if i > 0:
y_idx = torch.cat(
[y_idx,
torch.full((e.item(), ), i).long()])
y_idx = y_idx.to(self.device)
data.protein = torch_geometric.utils.to_dense_batch(
data.protein, y_idx)[0]
output = self.model(data)
else:
output = self.model(data)
data = Batch.from_data_list(data).to(self.device)
loss = 0
for i in range(self.tasks_num):
y_pred = output
y_label = data.y
loss = self.criterion[i](y_pred, y_label)
probs = F.softmax(y_pred.detach().cpu(), dim=-1)
y_out = probs.argmax(dim=1).numpy()
y_out_all.extend(y_out)
y_pred = probs[:, 1].view(-1).numpy()
try:
y_label_list[i].extend(y_label.cpu().numpy())
y_pred_list[i].extend(y_pred)
except:
y_label_list[i] = []
y_pred_list[i] = []
y_label_list[i].extend(y_label.cpu().numpy())
y_pred_list[i].extend(y_pred)
losses.append(loss.item())
val_roc = [
metrics.roc_auc_score(y_label_list[i], y_pred_list[i])
for i in range(self.tasks_num)
]
val_loss = np.array(losses).mean()
val_precision = metrics.precision_score(y_label_list[0],
np.array(y_out_all))
val_recall = metrics.recall_score(y_label_list[0], np.array(y_out_all))
return val_loss, np.array(val_roc).mean(), val_precision, val_recall
def train(self):
self.log('Training start...')
early_stop_cnt = 0
for epoch in range(self.option['epochs']):
trn_loss, trn_roc, trn_pre, trn_recall = self.train_iterations()
val_loss, val_roc, val_pre, val_recall = self.valid_iterations()
test_loss, test_roc, test_pre, test_recall = self.valid_iterations(
mode='test')
self.scheduler.step(val_loss)
lr_cur = self.scheduler.optimizer.param_groups[0]['lr']
self.log(
'Epoch:{} {} trn_loss:{:.3f} trn_roc:{:.3f} trn_precision:{:.3f} trn_recall:{:.3f} lr_cur:{:.5f}'
.format(epoch, self.option['dataset'], trn_loss, trn_roc,
trn_pre, trn_recall, lr_cur),
with_time=True)
self.log(
'Epoch:{} {} val_loss:{:.3f} val_roc:{:.3f} val_precision:{:.3f} val_recall:{:.3f} lr_cur:{:.5f}'
.format(epoch, self.option['dataset'], val_loss, val_roc,
val_pre, val_recall, lr_cur),
with_time=True)
self.log(
'Epoch:{} {} test_loss:{:.3f} test_roc:{:.3f} test_precision:{:.3f} test_recall:{:.3f} lr_cur:{:.5f}'
.format(epoch, self.option['dataset'], test_loss, test_roc,
test_pre, test_recall, lr_cur),
with_time=True)
self.records['val_auc'].append(val_roc)
if val_roc == np.array(self.records['val_auc']).max():
self.save_model_and_records(epoch,
test_pre,
test_recall,
test_roc,
final_save=False)
early_stop_cnt = 0
else:
early_stop_cnt += 1
if self.option[
'early_stop_patience'] > 0 and early_stop_cnt > self.option[
'early_stop_patience']:
self.log('Early stop hitted!')
break
self.log(
'The best epoch is {}, test_presion: {}, test_recall: {}, test_auc: {}'
.format(self.records['best_epoch'],
self.records['best_test_precision'],
self.records['best_test_recall'],
self.records['best_test_auc']))
return self.records['best_test_auc']
def predict(self):
self.model.eval()
dataloader = self.test_dataloader
ret = []
with torch.no_grad():
for data in tqdm(dataloader):
if not self.option['parallel']:
data = data.to(self.device)
target_len = data.pr_len
y_idx = torch.zeros(target_len[0]).long()
for i, e in enumerate(target_len):
if i > 0:
y_idx = torch.cat(
[y_idx,
torch.full((e.item(), ), i).long()])
y_idx = y_idx.to(self.device)
data.protein = torch_geometric.utils.to_dense_batch(
data.protein, y_idx)[0]
output = self.model(data)
else:
output = self.model(data)
output = F.softmax(output.detach().cpu(),
dim=-1)[:, 1].view(-1).numpy()
ret.extend(output)
return ret
def save_model_and_records(self,
epoch,
test_pre,
test_recall,
test_auc=None,
final_save=False):
if final_save:
# self.save_loss_records()
file_name = 'best_model_{}.ckpt'.format(self.option['seed'])
else:
file_name = 'best_model_{}.ckpt'.format(self.option['seed'])
self.records['best_epoch'] = epoch
self.records['best_test_precision'] = test_pre
self.records['best_test_recall'] = test_recall
self.records['best_test_auc'] = test_auc
if not self.option['parallel']:
model_dic = self.model.state_dict()
else:
model_dic = self.model.module.state_dict()
with open(os.path.join(self.save_path, file_name), 'wb') as f:
torch.save(model_dic, f)
self.log('Model saved at epoch {}'.format(epoch))
def save_loss_records(self):
trn_record = pd.DataFrame(
self.records['trn_record'],
columns=['epoch', 'trn_loss', 'trn_auc', 'trn_acc', 'lr'])
val_record = pd.DataFrame(
self.records['val_record'],
columns=['epoch', 'val_loss', 'val_auc', 'val_acc', 'lr'])
ret = pd.DataFrame({
'epoch': trn_record['epoch'],
'trn_loss': trn_record['trn_loss'],
'val_loss': val_record['val_loss'],
'trn_auc': trn_record['trn_auc'],
'val_auc': val_record['val_auc'],
'trn_lr': trn_record['lr'],
'val_lr': val_record['lr']
})
ret.to_csv(self.save_path + '/record.csv')
return ret
def load_best_ckpt(self):
ckpt_path = self.save_path + '/' + self.records['best_ckpt']
self.log('The best ckpt is {}'.format(ckpt_path))
self.load_ckpt(ckpt_path)
def load_ckpt(self, ckpt_path):
self.log('Ckpt loading: {}'.format(ckpt_path))
ckpt = torch.load(ckpt_path)
self.option = ckpt['option']
self.records = ckpt['records']
self.model.load_state_dict(ckpt['model_state_dict'])
def log(self, msg=None, msgs=None, with_time=False, show=True):
if with_time:
msg = msg + ' time elapsed {:.2f} hrs ({:.1f} mins)'.format(
(time.time() - self.start) / 3600.,
(time.time() - self.start) / 60.)
with open(self.save_path + '/log.txt', 'a+') as f:
if msgs:
self.log('#' * 80)
if '\n' not in msgs[0]: msgs = [m + '\n' for m in msgs]
f.writelines(msgs)
if show:
for x in msgs:
print(x, end='')
self.log('#' * 80)
if msg:
f.write(msg + '\n')
if show:
print(msg)
class Trainer():
def __init__(self, option, model, train_dataset, valid_dataset, test_dataset=None, weight=[[1.0, 1.0]],
tasks_num=17):
# Most important variable
self.option = option
self.device = torch.device("cuda:{}".format(option['gpu'][0]) if torch.cuda.is_available() else "cpu")
self.model = DataParallel(model).to(self.device) if option['parallel'] else model.to(self.device)
# Setting the train valid and test data loader
if self.option['parallel']:
self.train_dataloader = DataListLoader(train_dataset, batch_size=self.option['batch_size'], shuffle=True)
self.valid_dataloader = DataListLoader(valid_dataset, batch_size=self.option['batch_size'])
if test_dataset: self.test_dataloader = DataListLoader(test_dataset, batch_size=self.option['batch_size'])
else:
self.train_dataloader = DataLoader(train_dataset, batch_size=self.option['batch_size'], shuffle=True)
self.valid_dataloader = DataLoader(valid_dataset, batch_size=self.option['batch_size'])
if test_dataset: self.test_dataloader = DataLoader(test_dataset, batch_size=self.option['batch_size'])
self.save_path = self.option['exp_path']
# Setting the Adam optimizer with hyper-param
if option['focalloss']:
self.log('Using FocalLoss')
self.criterion = [FocalLoss(alpha=1 / w[0]) for w in weight] # alpha 0.965
else:
self.criterion = [torch.nn.CrossEntropyLoss(torch.Tensor(w).to(self.device), reduction='mean') for w in
weight]
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.option['lr'],
weight_decay=option['weight_decay'])
self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, mode='min', factor=0.7,
patience=self.option['lr_scheduler_patience'], min_lr=1e-6
)
# other
self.start = time.time()
self.tasks_num = tasks_num
self.records = {'trn_record': [], 'val_record': [], 'val_losses': [],
'best_ckpt': None, 'val_roc': [], 'val_prc': []}
self.log(msgs=['\t{}:{}\n'.format(k, v) for k, v in self.option.items()], show=False)
self.log('train set num:{} valid set num:{} test set num: {}'.format(
len(train_dataset), len(valid_dataset), len(test_dataset)))
self.log("total parameters:" + str(sum([p.nelement() for p in self.model.parameters()])))
self.log(msgs=str(model).split('\n'), show=False)
def train_iterations(self):
self.model.train()
losses = []
y_pred_list = {}
y_label_list = {}
for data in tqdm(self.train_dataloader):
self.optimizer.zero_grad()
data = data.to(self.device)
output = self.model(data)
loss = 0
for i in range(self.tasks_num):
y_pred = output[:, i * 2:(i + 1) * 2]
y_label = data.y[:, i].squeeze()
validId = np.where((y_label.cpu().numpy() == 0) | (y_label.cpu().numpy() == 1))[0]
if len(validId) == 0:
continue
if y_label.dim() == 0:
y_label = y_label.unsqueeze(0)
y_pred = y_pred[torch.tensor(validId).to(self.device)]
y_label = y_label[torch.tensor(validId).to(self.device)]
loss += self.criterion[i](y_pred, y_label)
y_pred = F.softmax(y_pred.detach().cpu(), dim=-1)[:, 1].view(-1).numpy()
try:
y_label_list[i].extend(y_label.cpu().numpy())
y_pred_list[i].extend(y_pred)
except:
y_label_list[i] = []
y_pred_list[i] = []
y_label_list[i].extend(y_label.cpu().numpy())
y_pred_list[i].extend(y_pred)
loss.backward()
# for name,parms in self.model.named_parameters():
# print('-->name:', name, '-->grad_requirs:', parms.requires_grad, \
# ' -->grad_value:', parms.grad)
self.optimizer.step()
losses.append(loss.item())
trn_roc = [metrics.roc_auc_score(y_label_list[i], y_pred_list[i]) for i in range(self.tasks_num)]
trn_prc = [metrics.auc(precision_recall_curve(y_label_list[i], y_pred_list[i])[1],
precision_recall_curve(y_label_list[i], y_pred_list[i])[0]) for i in
range(self.tasks_num)]
trn_loss = np.array(losses).mean()
return trn_loss, np.array(trn_roc).mean(), np.array(trn_prc).mean()
def valid_iterations(self, mode='valid'):
self.model.eval()
if mode == 'test' or mode == 'eval': dataloader = self.test_dataloader
if mode == 'valid': dataloader = self.valid_dataloader
losses = []
y_pred_list = {}
y_label_list = {}
with torch.no_grad():
for data in tqdm(dataloader):
data = data.to(self.device)
output = self.model(data)
loss = 0
for i in range(self.tasks_num):
y_pred = output[:, i * 2:(i + 1) * 2]
y_label = data.y[:, i].squeeze()
validId = np.where((y_label.cpu().numpy() == 0) | (y_label.cpu().numpy() == 1))[0]
if len(validId) == 0:
continue
if y_label.dim() == 0:
y_label = y_label.unsqueeze(0)
# print(len(validId))
# try:
# print(len(y_label))
# except:
# print(y_label)
# print(data.y[:,i])
# print(data.y)
y_pred = y_pred[torch.tensor(validId).to(self.device)]
y_label = y_label[torch.tensor(validId).to(self.device)]
loss = self.criterion[i](y_pred, y_label)
y_pred = F.softmax(y_pred.detach().cpu(), dim=-1)[:, 1].view(-1).numpy()
try:
y_label_list[i].extend(y_label.cpu().numpy())
y_pred_list[i].extend(y_pred)
except:
y_label_list[i] = []
y_pred_list[i] = []
y_label_list[i].extend(y_label.cpu().numpy())
y_pred_list[i].extend(y_pred)
losses.append(loss.item())
val_roc = [metrics.roc_auc_score(y_label_list[i], y_pred_list[i]) for i in range(self.tasks_num)]
val_prc = [metrics.auc(precision_recall_curve(y_label_list[i], y_pred_list[i])[1],
precision_recall_curve(y_label_list[i], y_pred_list[i])[0]) for i in
range(self.tasks_num)]
val_loss = np.array(losses).mean()
if mode == 'eval':
self.log('SEED {} DATASET {} The best test_loss:{:.3f} test_roc:{:.3f} test_prc:{:.3f}.'
.format(self.option['seed'], self.option['dataset'], val_loss, np.array(val_roc).mean(),
np.array(val_prc).mean()))
return val_loss, np.array(val_roc).mean(), np.array(val_prc).mean()
def train(self):
self.log('Training start...')
early_stop_cnt = 0
for epoch in range(self.option['epochs']):
trn_loss, trn_roc, trn_prc = self.train_iterations()
val_loss, val_roc, val_prc = self.valid_iterations()
test_loss, test_roc, test_prc = self.valid_iterations(mode='test')
self.scheduler.step(val_loss)
lr_cur = self.scheduler.optimizer.param_groups[0]['lr']
self.log('Epoch:{} {} trn_loss:{:.3f} trn_roc:{:.3f} trn_prc:{:.3f} lr_cur:{:.5f}'.
format(epoch, self.option['dataset'], trn_loss, trn_roc, trn_prc, lr_cur),
with_time=True)
self.log('Epoch:{} {} val_loss:{:.3f} val_roc:{:.3f} val_prc:{:.3f} lr_cur:{:.5f}'.
format(epoch, self.option['dataset'], val_loss, val_roc, val_prc, lr_cur),
with_time=True)
self.log('Epoch:{} {} test_loss:{:.3f} test_roc:{:.3f} test_prc:{:.3f} lr_cur:{:.5f}'.
format(epoch, self.option['dataset'], test_loss, test_roc, test_prc, lr_cur),
with_time=True)
self.records['val_roc'].append(val_roc)
self.records['val_prc'].append(val_prc)
self.records['val_record'].append([epoch, val_loss, val_roc, val_prc, lr_cur])
self.records['trn_record'].append([epoch, trn_loss, trn_roc, trn_prc, lr_cur])
if val_roc == np.array(self.records['val_roc']).max() or val_prc == np.array(self.records['val_prc']).max():
self.save_model_and_records(epoch)
early_stop_cnt = 0
else:
early_stop_cnt += 1
if self.option['early_stop_patience'] > 0 and early_stop_cnt > self.option['early_stop_patience']:
self.log('Early stop hitted!')
break
self.save_model_and_records(epoch, final_save=True)
def save_model_and_records(self, epoch, final_save=False):
if final_save:
self.save_loss_records()
file_name = 'best_model.ckpt'
else:
file_name = 'best_model.ckpt'
self.records['best_ckpt'] = file_name
with open(os.path.join(self.save_path, file_name), 'wb') as f:
torch.save({
'option': self.option,
'records': self.records,
'model_state_dict': self.model.state_dict(),
}, f)
self.log('Model saved at epoch {}'.format(epoch))
def save_loss_records(self):
trn_record = pd.DataFrame(self.records['trn_record'],
columns=['epoch', 'trn_loss', 'trn_auc', 'trn_acc', 'lr'])
val_record = pd.DataFrame(self.records['val_record'],
columns=['epoch', 'val_loss', 'val_auc', 'val_acc', 'lr'])
ret = pd.DataFrame({
'epoch': trn_record['epoch'],
'trn_loss': trn_record['trn_loss'],
'val_loss': val_record['val_loss'],
'trn_auc': trn_record['trn_auc'],
'val_auc': val_record['val_auc'],
'trn_lr': trn_record['lr'],
'val_lr': val_record['lr']
})
ret.to_csv(self.save_path + '/record.csv')
return ret
def load_best_ckpt(self):
ckpt_path = self.save_path + '/' + self.records['best_ckpt']
self.log('The best ckpt is {}'.format(ckpt_path))
self.load_ckpt(ckpt_path)
def load_ckpt(self, ckpt_path):
self.log('Ckpt loading: {}'.format(ckpt_path))
ckpt = torch.load(ckpt_path)
self.option = ckpt['option']
self.records = ckpt['records']
self.model.load_state_dict(ckpt['model_state_dict'])
def log(self, msg=None, msgs=None, with_time=False, show=True):
if with_time: msg = msg + ' time elapsed {:.2f} hrs ({:.1f} mins)'.format(
(time.time() - self.start) / 3600.,
(time.time() - self.start) / 60.
)
with open(self.save_path + '/log.txt', 'a+') as f:
if msgs:
self.log('#' * 80)
if '\n' not in msgs[0]: msgs = [m + '\n' for m in msgs]
f.writelines(msgs)
if show:
for x in msgs:
print(x, end='')
self.log('#' * 80)
if msg:
f.write(msg + '\n')
if show:
print(msg)
def train_cross_validation(model_cls,
dataset,
num_clusters,
dropout=0.0,
lr=1e-4,
weight_decay=1e-2,
num_epochs=200,
n_splits=10,
use_gpu=True,
dp=False,
ddp=True,
comment='',
tb_service_loc='192.168.192.57:6006',
batch_size=1,
num_workers=0,
pin_memory=False,
cuda_device=None,
fold_no=None,
saved_model_path=None,
device_ids=None,
patience=50,
seed=None,
save_model=True,
c_reg=0,
base_log_dir='runs',
base_model_save_dir='saved_models'):
"""
:param c_reg:
:param save_model: bool
:param seed:
:param patience: for early stopping
:param device_ids: for ddp
:param saved_model_path:
:param fold_no:
:param ddp: DDP
:param cuda_device:
:param pin_memory: DataLoader args https://devblogs.nvidia.com/how-optimize-data-transfers-cuda-cc/
:param num_workers: DataLoader args
:param model_cls: pytorch Module cls
:param dataset: pytorch Dataset cls
:param dropout:
:param lr:
:param weight_decay:
:param num_epochs:
:param n_splits: number of kFolds
:param use_gpu: bool
:param dp: bool
:param comment: comment in the logs, to filter runs in tensorboard
:param tb_service_loc: tensorboard service location
:param batch_size: Dataset args not DataLoader
:return:
"""
saved_args = locals()
seed = int(time.time() % 1e4 * 1e5) if seed is None else seed
saved_args['random_seed'] = seed
torch.manual_seed(seed)
np.random.seed(seed)
if use_gpu:
torch.cuda.manual_seed_all(seed)
if ddp and not torch.distributed.is_initialized(): # initialize ddp
dist.init_process_group('nccl',
init_method='tcp://localhost:{}'.format(
find_open_port()),
world_size=1,
rank=0)
model_name = model_cls.__name__
if not cuda_device:
if device_ids and (ddp or dp):
device = device_ids[0]
else:
device = torch.device(
'cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
else:
device = cuda_device
device_count = torch.cuda.device_count() if dp else 1
device_count = len(device_ids) if (device_ids is not None and
(dp or ddp)) else device_count
if device_count > 1:
print("Let's use", device_count, "GPUs!")
# batch_size = batch_size * device_count
log_dir_base = get_model_log_dir(comment, model_name)
if tb_service_loc is not None:
print("TensorBoard available at http://{1}/#scalars®exInput={0}".
format(log_dir_base, tb_service_loc))
else:
print("Please set up TensorBoard")
criterion = nn.CrossEntropyLoss()
# get test set
folds = StratifiedKFold(n_splits=n_splits, shuffle=False)
train_val_idx, test_idx = list(
folds.split(np.zeros(len(dataset)), dataset.data.y.numpy()))[0]
test_dataset = dataset.__indexing__(test_idx)
train_val_dataset = dataset.__indexing__(train_val_idx)
print("Training {0} {1} models for cross validation...".format(
n_splits, model_name))
# folds, fold = KFold(n_splits=n_splits, shuffle=False, random_state=seed), 0
folds = StratifiedKFold(n_splits=n_splits, shuffle=False)
iter = folds.split(np.zeros(len(train_val_dataset)),
train_val_dataset.data.y.numpy())
fold = 0
for train_idx, val_idx in tqdm_notebook(iter, desc='CV', leave=False):
fold += 1
if fold_no is not None:
if fold != fold_no:
continue
writer = SummaryWriter(log_dir=osp.join(base_log_dir, log_dir_base +
str(fold)))
model_save_dir = osp.join(base_model_save_dir,
log_dir_base + str(fold))
print("creating dataloader tor fold {}".format(fold))
model = model_cls(writer,
num_clusters=num_clusters,
in_dim=dataset.data.x.shape[1],
out_dim=int(dataset.data.y.max() + 1),
dropout=dropout)
# My Batch
train_dataset = train_val_dataset.__indexing__(train_idx)
val_dataset = train_val_dataset.__indexing__(val_idx)
train_dataset = dataset_gather(
train_dataset,
seed=0,
n_repeat=1,
n_splits=int(len(train_dataset) / batch_size) + 1)
val_dataset = dataset_gather(
val_dataset,
seed=0,
n_repeat=1,
n_splits=int(len(val_dataset) / batch_size) + 1)
train_dataloader = DataLoader(train_dataset,
shuffle=True,
batch_size=device_count,
collate_fn=lambda data_list: data_list,
num_workers=num_workers,
pin_memory=pin_memory)
val_dataloader = DataLoader(val_dataset,
shuffle=False,
batch_size=device_count,
collate_fn=lambda data_list: data_list,
num_workers=num_workers,
pin_memory=pin_memory)
# if fold == 1 or fold_no is not None:
print(model)
writer.add_text('model_summary', model.__repr__())
writer.add_text('training_args', str(saved_args))
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=weight_decay,
amsgrad=False)
# optimizer = torch.optim.SGD(model.parameters(), lr=lr, weight_decay=weight_decay)
if ddp:
model = model.cuda() if device_ids is None else model.to(
device_ids[0])
model = nn.parallel.DistributedDataParallel(model,
device_ids=device_ids)
elif dp and use_gpu:
model = model.cuda() if device_ids is None else model.to(
device_ids[0])
model = DataParallel(model, device_ids=device_ids)
elif use_gpu:
model = model.to(device)
if saved_model_path is not None:
model.load_state_dict(torch.load(saved_model_path))
best_map, patience_counter, best_score = 0.0, 0, -np.inf
for epoch in tqdm_notebook(range(1, num_epochs + 1),
desc='Epoch',
leave=False):
for phase in ['train', 'validation']:
if phase == 'train':
model.train()
dataloader = train_dataloader
else:
model.eval()
dataloader = val_dataloader
# Logging
running_total_loss = 0.0
running_corrects = 0
running_reg_loss = 0.0
running_nll_loss = 0.0
epoch_yhat_0, epoch_yhat_1 = torch.tensor([]), torch.tensor([])
epoch_label, epoch_predicted = torch.tensor([]), torch.tensor(
[])
for data_list in tqdm_notebook(dataloader,
desc=phase,
leave=False):
# TODO: check devices
if dp:
data_list = to_cuda(data_list,
(device_ids[0] if device_ids
is not None else 'cuda'))
y_hat, reg = model(data_list)
# y_hat = y_hat.reshape(batch_size, -1)
y = torch.tensor([],
dtype=dataset.data.y.dtype,
device=device)
for data in data_list:
y = torch.cat([y, data.y.view(-1).to(device)])
loss = criterion(y_hat, y)
reg_loss = -reg
total_loss = (loss + reg_loss * c_reg).sum()
if phase == 'train':
# print(torch.autograd.grad(y_hat.sum(), model.saved_x, retain_graph=True))
optimizer.zero_grad()
total_loss.backward(retain_graph=True)
nn.utils.clip_grad_norm_(model.parameters(), 2.0)
optimizer.step()
_, predicted = torch.max(y_hat, 1)
label = y
running_nll_loss += loss.item()
running_total_loss += total_loss.item()
running_reg_loss += reg.sum().item()
running_corrects += (predicted == label).sum().item()
epoch_yhat_0 = torch.cat(
[epoch_yhat_0, y_hat[:, 0].detach().view(-1).cpu()])
epoch_yhat_1 = torch.cat(
[epoch_yhat_1, y_hat[:, 1].detach().view(-1).cpu()])
epoch_label = torch.cat(
[epoch_label,
label.detach().cpu().float()])
epoch_predicted = torch.cat(
[epoch_predicted,
predicted.detach().cpu().float()])
# precision = sklearn.metrics.precision_score(epoch_label, epoch_predicted, average='micro')
# recall = sklearn.metrics.recall_score(epoch_label, epoch_predicted, average='micro')
# f1_score = sklearn.metrics.f1_score(epoch_label, epoch_predicted, average='micro')
accuracy = sklearn.metrics.accuracy_score(
epoch_label, epoch_predicted)
epoch_total_loss = running_total_loss / dataloader.__len__()
epoch_nll_loss = running_nll_loss / dataloader.__len__()
epoch_reg_loss = running_reg_loss / dataloader.dataset.__len__(
)
writer.add_scalars(
'nll_loss', {'{}_nll_loss'.format(phase): epoch_nll_loss},
epoch)
writer.add_scalars('accuracy',
{'{}_accuracy'.format(phase): accuracy},
epoch)
# writer.add_scalars('{}_APRF'.format(phase),
# {
# 'accuracy': accuracy,
# 'precision': precision,
# 'recall': recall,
# 'f1_score': f1_score
# },
# epoch)
if epoch_reg_loss != 0:
writer.add_scalars(
'reg_loss'.format(phase),
{'{}_reg_loss'.format(phase): epoch_reg_loss}, epoch)
# writer.add_histogram('hist/{}_yhat_0'.format(phase),
# epoch_yhat_0,
# epoch)
# writer.add_histogram('hist/{}_yhat_1'.format(phase),
# epoch_yhat_1,
# epoch)
# Save Model & Early Stopping
if phase == 'validation':
model_save_path = model_save_dir + '-{}-{}-{:.3f}-{:.3f}'.format(
model_name, epoch, accuracy, epoch_nll_loss)
if accuracy > best_map:
best_map = accuracy
model_save_path = model_save_path + '-best'
score = -epoch_nll_loss
if score > best_score:
patience_counter = 0
best_score = score
else:
patience_counter += 1
# skip 10 epoch
# best_score = best_score if epoch > 10 else -np.inf
if save_model:
for th, pfix in zip(
[0.8, 0.75, 0.7, 0.5, 0.0],
['-perfect', '-great', '-good', '-bad', '-miss']):
if accuracy >= th:
model_save_path += pfix
break
if epoch > 10:
torch.save(model.state_dict(), model_save_path)
writer.add_scalars('best_val_accuracy',
{'{}_accuracy'.format(phase): best_map},
epoch)
writer.add_scalars(
'best_nll_loss',
{'{}_nll_loss'.format(phase): -best_score}, epoch)
if patience_counter >= patience:
print("Stopped at epoch {}".format(epoch))
return
print("Done !")
def train_cummunity_detection(model_cls,
dataset,
dropout=0.0,
lr=1e-3,
weight_decay=1e-2,
num_epochs=200,
n_splits=10,
use_gpu=True,
dp=False,
ddp=False,
comment='',
tb_service_loc='192.168.192.57:6006',
batch_size=1,
num_workers=0,
pin_memory=False,
cuda_device=None,
ddp_port='23456',
fold_no=None,
device_ids=None,
patience=20,
seed=None,
save_model=False,
supervised=False):
"""
:param save_model: bool
:param seed:
:param patience: for early stopping
:param device_ids: for ddp
:param saved_model_path:
:param fold_no:
:param ddp_port:
:param ddp: DDP
:param cuda_device:
:param pin_memory: DataLoader args https://devblogs.nvidia.com/how-optimize-data-transfers-cuda-cc/
:param num_workers: DataLoader args
:param model_cls: pytorch Module cls
:param dataset: pytorch Dataset cls
:param dropout:
:param lr:
:param weight_decay:
:param num_epochs:
:param n_splits: number of kFolds
:param use_gpu: bool
:param dp: bool
:param comment: comment in the logs, to filter runs in tensorboard
:param tb_service_loc: tensorboard service location
:param batch_size: Dataset args not DataLoader
:return:
"""
saved_args = locals()
seed = int(time.time() % 1e4 * 1e5) if seed is None else seed
saved_args['random_seed'] = seed
torch.manual_seed(seed)
np.random.seed(seed)
if use_gpu:
torch.cuda.manual_seed_all(seed)
if ddp and not torch.distributed.is_initialized(): # initialize ddp
dist.init_process_group(
'nccl',
init_method='tcp://localhost:{}'.format(ddp_port),
world_size=1,
rank=0)
model_name = model_cls.__name__
if not cuda_device:
if device_ids and (ddp or dp):
device = device_ids[0]
else:
device = torch.device(
'cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
else:
device = cuda_device
device_count = torch.cuda.device_count() if dp else 1
device_count = len(device_ids) if (device_ids is not None and
(dp or ddp)) else device_count
if device_count > 1:
print("Let's use", device_count, "GPUs!")
# batch_size = batch_size * device_count
log_dir_base = get_model_log_dir(comment, model_name)
if tb_service_loc is not None:
print("TensorBoard available at http://{1}/#scalars®exInput={0}".
format(log_dir_base, tb_service_loc))
else:
print("Please set up TensorBoard")
print("Training {0} {1} models for cross validation...".format(
n_splits, model_name))
folds, fold = KFold(n_splits=n_splits, shuffle=False, random_state=seed), 0
print(dataset.__len__())
for train_idx, test_idx in tqdm_notebook(folds.split(
list(range(dataset.__len__())), list(range(dataset.__len__()))),
desc='models',
leave=False):
fold += 1
if fold_no is not None:
if fold != fold_no:
continue
writer = SummaryWriter(log_dir=osp.join('runs', log_dir_base +
str(fold)))
model_save_dir = osp.join('saved_models', log_dir_base + str(fold))
print("creating dataloader tor fold {}".format(fold))
model = model_cls(writer, dropout=dropout)
# My Batch
train_dataset = dataset.__indexing__(train_idx)
test_dataset = dataset.__indexing__(test_idx)
train_dataset = dataset_gather(train_dataset,
n_repeat=1,
n_splits=int(
len(train_dataset) / batch_size))
train_dataloader = DataLoader(train_dataset,
shuffle=True,
batch_size=device_count,
collate_fn=lambda data_list: data_list,
num_workers=num_workers,
pin_memory=pin_memory)
test_dataloader = DataLoader(test_dataset,
shuffle=False,
batch_size=device_count,
collate_fn=lambda data_list: data_list,
num_workers=num_workers,
pin_memory=pin_memory)
print(model)
writer.add_text('model_summary', model.__repr__())
writer.add_text('training_args', str(saved_args))
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=weight_decay,
amsgrad=False)
# optimizer = torch.optim.SGD(model.parameters(), lr=lr, weight_decay=weight_decay)
if ddp:
model = model.cuda() if device_ids is None else model.to(
device_ids[0])
model = nn.parallel.DistributedDataParallel(model,
device_ids=device_ids)
elif dp and use_gpu:
model = model.cuda() if device_ids is None else model.to(
device_ids[0])
model = DataParallel(model, device_ids=device_ids)
elif use_gpu:
model = model.to(device)
for epoch in tqdm_notebook(range(1, num_epochs + 1),
desc='Epoch',
leave=False):
for phase in ['train', 'validation']:
if phase == 'train':
model.train()
dataloader = train_dataloader
else:
model.eval()
dataloader = test_dataloader
# Logging
running_total_loss = 0.0
running_reg_loss = 0.0
running_overlap = 0.0
for data_list in tqdm_notebook(dataloader,
desc=phase,
leave=False):
# TODO: check devices
if dp:
data_list = to_cuda(data_list,
(device_ids[0] if device_ids
is not None else 'cuda'))
y_hat, reg = model(data_list)
y = torch.tensor([],
dtype=dataset.data.y.dtype,
device=device)
for data in data_list:
y = torch.cat([y, data.y.view(-1).to(device)])
if supervised:
loss = permutation_invariant_loss(y_hat, y)
# criterion = nn.NLLLoss()
# loss = criterion(y_hat, y)
else:
loss = -reg
total_loss = loss
if phase == 'train':
# print(torch.autograd.grad(y_hat.sum(), model.saved_x, retain_graph=True))
optimizer.zero_grad()
total_loss.backward(retain_graph=True)
nn.utils.clip_grad_norm_(model.parameters(), 2.0)
optimizer.step()
_, predicted = torch.max(y_hat, 1)
label = y
if supervised:
overlap_score = normalized_overlap(
label.int().cpu().numpy(),
predicted.int().cpu().numpy(), 0.25)
# overlap_score = overlap(label.int().cpu().numpy(), predicted.int().cpu().numpy())
running_overlap += overlap_score
print(reg, overlap_score, loss)
running_total_loss += total_loss.item()
running_reg_loss += reg.sum().item()
epoch_total_loss = running_total_loss / dataloader.__len__()
epoch_reg_loss = running_reg_loss / dataloader.dataset.__len__(
)
if supervised:
epoch_overlap = running_overlap / dataloader.__len__()
writer.add_scalars(
'overlap'.format(phase),
{'{}_overlap'.format(phase): epoch_overlap}, epoch)
writer.add_scalars(
'reg_loss'.format(phase),
{'{}_reg_loss'.format(phase): epoch_reg_loss}, epoch)
print("Done !")