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 init_model(model_cls,
log_dir_base,
fold_no,
device_ids=None,
use_gpu=False,
dp=False,
ddp=False,
tb_dir='runs',
lr=1e-3,
weight_decay=1e-2):
writer = SummaryWriter(log_dir=osp.join(tb_dir, log_dir_base))
model = model_cls(writer)
writer.add_text('model_summary', model.__repr__())
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_ids[0])
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
return model, optimizer, writer, device_count
def apply_dataparallel(model, cfgs):
return DataParallel(model)
self.lin1 = torch.nn.Linear(64, 128)
self.lin2 = torch.nn.Linear(128, dataset.num_classes)
def forward(self, data):
print('Inside Model: num graphs: {}, device: {}'.format(
data.num_graphs, data.batch.device))
x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
x = F.elu(self.conv1(x, edge_index, edge_attr))
x = F.elu(self.conv2(x, edge_index, edge_attr))
x = global_mean_pool(x, data.batch)
x = F.elu(self.lin1(x))
return F.log_softmax(self.lin2(x), dim=1)
model = Net()
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
model = DataParallel(model)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for data_list in loader:
optimizer.zero_grad()
output = model(data_list)
print('Outside Model: num graphs: {}'.format(output.size(0)))
y = torch.cat([data.y for data in data_list]).to(output.device)
loss = F.nll_loss(output, y)
loss.backward()
optimizer.step()
id_euc = sphg(pos, 0.6, batch=batch, max_num_neighbors=16)
x9 = self.conv3(x8, pos, B, N, id_euc)
x10 = self.lin3(x9.view(B, N, -1))
x = x10.max(1)[0] # [B, C]
return self.fc(x)
# Train and test
model = Net(train_dataset.num_classes)
model = model.to(device)
model.load_state_dict(torch.load('weight.pth',
map_location=f'cuda:{device_list[0]}'),
strict=True)
if cuda: model = DataParallel(model, device_ids=device_list)
optimizer = torch.optim.Adam([{
'params': model.parameters(),
'initial_lr': base_lr
}],
lr=base_lr,
weight_decay=1e-4)
# optimizer.load_state_dict(torch.load('geocnn_optimizer.pt', map_location=f'cuda:{device_list[0]}').state_dict())
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
epoch,
eta_min=0.00001,
last_epoch=-1)
criterion = cal_loss
def train(epoch):
y = torch.cat([data.y for data in data]).to(output.device)
loss = F.nll_loss(output, y)
epoch_test_loss += loss.detach().item()
pred = output.max(dim=1)[1]
correct += pred.eq(y).sum().item()
return epoch_test_loss / len(test_dataset), correct / len(test_dataset)
t0 = time.time()
num_features = train_dataset.num_features
n_classes = train_dataset.num_classes
#model = GCN_PYG(num_features, 96, n_classes, 0)
model = GAT_PYG(num_features, n_classes, 32, 8, 0.5)
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
model = DataParallel(model, [1])
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay = 0)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',
factor=0.5,
patience=25,
verbose=True)
dur = []
for epoch in range(1, 201):
t1 = time.time()
train_loss, train_acc, optimizer= train_pyg('gcn', model, train_loader, device, optimizer)
dur.append(time.time() - t1)
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 __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 __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)
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 !")
def TrainingNet(dataset,modelName,params,num_pre_epochs,num_epochs,NumCutoff,optimizerName,LinkPredictionMethod,MonteSize,savepath):
Batch_size=params[0]
VectorPairs=params[4]
StartTopoCoeffi=params[5]
WeightCorrectionCoeffi=params[6]
interval=params[7]
root='/git/data/GraphData/'+dataset
TestAccs=[]
for Monte_iter in range(MonteSize):
# Data
NewNetworkSizeAdjust=[]
WeightsDynamicsEvolution=[]
trainValRatio=[0.2,0.4]
# model
if dataset=='Cora' or dataset =='Citeseer' or dataset =='Pubmed':
datasetroot= Planetoid(root=root, name=dataset,transform =T.NormalizeFeatures()).shuffle()
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
""" train_mask, val_mask,test_mask=DataSampler(trainValRatio,datasetroot.data.num_nodes)
DataMask={}
DataMask['train_mask']=train_mask
DataMask['val_mask']=val_mask
DataMask['test_mask']=test_mask
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)"""
num_features=datasetroot.num_features
num_classes=datasetroot.num_classes
criterion = nn.CrossEntropyLoss()
elif dataset =="CoraFull":
datasetroot = CoraFull(root=root,transform =T.NormalizeFeatures()).shuffle()
"""train_mask, val_mask,test_mask=DataSampler(trainValRatio,datasetroot.data.num_nodes)
DataMask={}
DataMask['train_mask']=train_mask
DataMask['val_mask']=val_mask
DataMask['test_mask']=test_mask"""
criterion = nn.CrossEntropyLoss()
num_features=datasetroot.num_features
num_classes=datasetroot.num_classes
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=False)
elif dataset=='ENZYMES' or dataset=='MUTAG':
datasetroot=TUDataset(root,name=dataset,use_node_attr=True)
trainloader = DataLoader(datasetroot, batch_size=Batch_size, shuffle=True)
num_features=datasetroot.num_features
num_classes=datasetroot.num_classes
elif dataset =="PPI":
train_dataset = PPI(root, split='train')
val_dataset = PPI(root, split='val')
test_dataset = PPI(root, split='test')
trainloader = DataListLoader(train_dataset, batch_size=Batch_size, shuffle=True)
valloader = DataListLoader(val_dataset, batch_size=100, shuffle=False)
testloader = DataListLoader(test_dataset, batch_size=100, shuffle=False)
num_classes=train_dataset.num_classes
num_features=train_dataset.num_features
criterion = torch.nn.BCEWithLogitsLoss()
elif dataset =="Reddit":
datasetroot=Reddit(root)
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
testloader = DataListLoader(datasetroot, batch_size=2, shuffle=False)
elif dataset=="Amazon":
datasetroot=Amazon(root, "Photo", transform=None, pre_transform=None)
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False)
elif dataset=='MNIST':
datasetroot = MNISTSuperpixels(root=root, transform=T.Cartesian())
trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False)
elif dataset=='CIFAR10':
pass
else:
raise Exception("Input wrong datatset!!")
width=ContractionLayerCoefficients(num_features,*params[1:3])
net =ChooseModel(modelName,num_features,num_classes,width)
FileName="{}-{}-param_{}_{}_{}_{}-monte_{}".format(dataset,modelName,interval,WeightCorrectionCoeffi,StartTopoCoeffi,VectorPairs,Monte_iter)
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
criterion = criterion.to(device)
net = DataParallel(net)
net = net.to(device)
optimizer = getattr(optim,optimizerName)(net.parameters(), lr=params[3], momentum=0.9, weight_decay=5e-4)
model_to_save='./checkpoint/{}-{}-param_{}_{}_{}_{}-ckpt.pth'.format(dataset,modelName,params[0],params[1],params[5],params[4])
if resume=="True" and os.path.exists(model_to_save):
[net,optimizer,TrainConvergence,TestConvergence,Acc]=ResumeModel(net,optimizer,model_to_save)
start_epoch=len(TrainConvergence)
else:
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
#cudnn.benchmark = True
logging('dataset:{}, Batch size: {}, Number of layers:{} ConCoeff: {}, LR:{}, MonteSize:{}'.format(dataset, params[0], params[1],params[2],params[3],Monte_iter))
mark="{}{}Convergence/DiagElement-{}".format(savepath,dataset,FileName)
markweights="{}{}Convergence/WeightChanges-{}".format(savepath,dataset,FileName)
PreTrainConvergence,PreTestConvergence,PreAcc=TrainPart(start_epoch,num_pre_epochs,num_classes, trainloader,net,optimizer,criterion,NumCutoff,LinkPredictionMethod,VectorPairs,WeightCorrectionCoeffi,StartTopoCoeffi,mark,markweights,model_to_save,False)
print('dataset: {}, model name:{}, epoch:{},Pre-train error:{}; Pre-test error:{}; test acc:{}'.format(dataset,modelName,num_pre_epochs,PreTrainConvergence[-1],PreTestConvergence[-1],PreAcc))
NewNetworksize=RetainNetworkSize(net,params[2])
OptimizedNet=ChooseModel(modelName,num_features,num_classes,NewNetworksize[0:-1])
NewNetworksize.insert(0,num_features)
NewNetworkSizeAdjust.append(NewNetworksize[0:-1])
print(NewNetworkSizeAdjust)
#OptimizedNet.apply(init_weights)
OptimizedNet = DataParallel(OptimizedNet)
OptimizedNet = OptimizedNet.to(device)
cudnn.benchmark = True
# Begin Pre training
if optimizerName =="SGD":
optimizerNew = getattr(optim,optimizerName)(OptimizedNet.parameters(), lr=params[3], momentum=0.9, weight_decay=5e-4)
elif optimizerName =="Adam":
optimizerNew = getattr(optim,optimizerName)(OptimizedNet.parameters(), lr=params[3], betas=(0.9, 0.999), eps=1e-08, weight_decay=5e-4, amsgrad=False)
TrainConvergence,TestConvergence,TestAcc=TrainPart(start_epoch,num_epochs,datasetroot.num_classes, trainloader,OptimizedNet,optimizerNew,criterion,
NumCutoff,LinkPredictionMethod,VectorPairs,WeightCorrectionCoeffi,StartTopoCoeffi,mark,markweights,model_to_save,True)
np.save("{}/{}Convergence/AlgebraicConectivityTrainConvergence-{}".format(savepath,dataset,FileName),TrainConvergence)
np.save("{}/{}Convergence/AlgebraicConectivityTestConvergence-{}".format(savepath,dataset,FileName),TestConvergence)
#np.save("{}/{}Convergence/NewNetworkSizeAdjust-{}".format(savepath,dataset,FileName),NewNetworkSizeAdjust)
#torch.cuda.empty_cache()
print('dataset: {}, model name:{}, resized network size: {}, the train error: {},test error: {}, test acc:{}\n'.format(dataset,modelName,NewNetworksize[0:-1],num_epochs,TrainConvergence[-1],TestConvergence[-1],TestAcc))
np.save("{}/{}Convergence/AlgebraicConectivityMeanTestAccs-{}".format(savepath,dataset,FileName),TestAccs.append(TestAcc))
TestAccs.append(TestAcc)
print_nvidia_useage()
3)).view(-1, 3)
x1 = self.conv1(pos, batch)
x2 = self.conv2(x1, batch)
x3 = self.conv3(x2, batch)
x4 = self.lin1(x3)
x5 = global_max_pool(x4, batch)
x6 = x5.repeat([1, 2048]).view(-1, 1024)
x7 = torch.cat([x2, x3, x6], dim=1)
out = self.mlp(x7)
return F.log_softmax(out, dim=1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net(train_dataset.num_classes, k=30)
model = DataParallel(model).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.8)
def train():
model.train()
total_loss = correct_nodes = total_nodes = 0
for i, data_list in enumerate(train_loader):
optimizer.zero_grad()
out = model(data_list)
y = torch.cat([data.y for data in data_list]).to(device)
loss = F.nll_loss(out, y)
loss.backward()
optimizer.step()
print('Importing structures.')
trainset = Structures(root='./datasets/{}_train/'.format(p.dataset),
prefix=p.dataset)
samples = len(trainset)
cutoff = int(np.floor(samples*(1-p.validation_split)))
validset = trainset[cutoff:]
trainset = trainset[:cutoff]
if p.shuffle_dataset:
trainset = trainset.shuffle()
n_features = trainset.get(0).x.shape[1]
print('Setting up model...')
model = p.model_type(6, heads=p.heads).to(device)
model = DataParallel(model).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learn_rate, weight_decay=p.weight_decay)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',
# factor=p.lr_decay,
# patience=p.patience)
writer = SummaryWriter(comment='model:{}_lr:{}_lr_decay:{}_shuffle:{}_seed:{}'.format(
p.version,
learn_rate,
p.lr_decay,
p.shuffle_dataset,
p.random_seed))
# axes = [0, 1, 2]
max_roc_auc = 0
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)
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,
)
def GCN(dataset, params, Epochs, MonteSize, width, lr, savepath):
Batch_size = int(params[0])
for Monte_iter in range(MonteSize):
# Data
best_loss = float('inf') # best test loss
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
TrainConvergence = []
TestConvergence = []
# model
root = '/data/GraphData/' + dataset
if dataset == 'Cora':
model_name = "GCN3"
datasetroot = Planetoid(root=root, name=dataset).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], Monte_iter)
if resume and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= Epochs - 1:
continue
else:
net = Net(datasetroot, width)
elif dataset == 'ENZYMES' or dataset == 'MUTAG':
model_name = "topk_pool_Net"
root = '/data/GraphData' + dataset
datasetroot = TUDataset(root, name=dataset)
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], Monte_iter)
if resume and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= Epochs - 1:
continue
else:
net = topk_pool_Net(datasetroot, width)
elif dataset == 'MNIST':
datasetroot = MNISTSuperpixels(root='/data/GraphData/' + dataset,
transform=T.Cartesian()).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_name = 'SPlineNet'
model_to_save = './checkpoint/{}-{}-param_{}_{}-Mon_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], Monte_iter)
if resume and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= Epochs - 1:
continue
else:
#net=Net(datasetroot,width)
net = SPlineNet(datasetroot, width)
elif dataset == 'CIFAR10':
if resume and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= Epochs - 1:
continue
else:
net = getattr(CIFAR10_resnet, 'Resnet20_CIFAR10')(params[1])
else:
raise Exception(
"The dataset is:{}, it isn't existed.".format(dataset))
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
torch.cuda.is_available()
net = DataParallel(net)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net = net.to(device)
#cudnn.benchmark = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=5e-4)
for epoch in range(start_epoch, start_epoch + Epochs):
if epoch < Epochs:
logging(
'Batch size: {},ConCoeff: {},MonteSize:{},epoch:{}'.format(
params[0], params[1], Monte_iter, epoch))
TrainLoss = train(trainloader, net, optimizer, criterion)
TrainConvergence.append(statistics.mean(TrainLoss))
TestConvergence.append(
statistics.mean(test(testloader, net, criterion)))
else:
break
if TestConvergence[epoch] < best_loss:
logging('Saving..')
state = {
'net': net.module,
'TrainConvergence': TrainConvergence,
'TestConvergence': TestConvergence,
'epoch': epoch,
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, model_to_save)
best_loss = TestConvergence[epoch]
if not os.path.exists('./%s' % model_name):
os.makedirs('./%s' % model_name)
torch.save(
net.module.state_dict(),
'./%s/%s_%s_%s_%s_%s_pretrain.pth' %
(model_name, dataset, model_name, params[0], params[1],
Epochs))
else:
pass
## save recurrence plots
if epoch % 20 == 0:
save_recurrencePlots_file = "../Results/RecurrencePlots/RecurrencePlots_{}_{}_BatchSize{}_ConCoeffi{}_epoch{}.png".format(
dataset, model_name, params[0], params[1], epoch)
save_recurrencePlots(net, save_recurrencePlots_file)
FileName = "{}-{}-param_{}_{}-monte_{}".format(dataset, model_name,
params[0], params[1],
Monte_iter)
np.save(savepath + 'TrainConvergence-' + FileName, TrainConvergence)
np.save(savepath + 'TestConvergence-' + FileName, TestConvergence)
torch.cuda.empty_cache()
print_nvidia_useage()
if return_output == True:
return TestConvergence[-1], net.module.fc.weight
else:
pass
parameters = list(model.parameters())
optimizer = torch.optim.Adam(params=parameters, lr=args.lr)
total_params = sum(p.numel() for param in parameters for p in param)
print(f'Total number of parameters is {total_params}')
if args.model == 'DGCNN':
print(f'SortPooling k is set to {model.k}')
log_file = os.path.join(args.res_dir, 'log.txt')
with open(log_file, 'a') as f:
print(f'Total number of parameters is {total_params}', file=f)
if args.model == 'DGCNN':
print(f'SortPooling k is set to {model.k}', file=f)
start_epoch = 1
if args.multi_gpu: model = DataParallel(model)
model = model.to(device)
if args.continue_from is not None:
model.load_state_dict(
torch.load(
os.path.join(
args.res_dir,
'model_checkpoint{}.pth'.format(args.continue_from))))
optimizer.load_state_dict(
torch.load(
os.path.join(
args.res_dir,
'optimizer_checkpoint{}.pth'.format(args.continue_from))))
start_epoch = args.continue_from + 1
args.epochs -= args.continue_from
def train_RGNN(tr_dataset, te_dataset, n_epochs, batch_size, lr, z_dim, K, dropout, adj_type, learn_edge, lambda1,
lambda2, domain_adaptation, lambda_dat, label_type, ckpt_save_name=None, ckpt_load=None):
# log hyper-parameter
logger.critical('batch_size {}, lr {}, z_dim {}, K {}, dropout {}, adj_type {}, learn_edge {}, lambda1 {},'
'lambda2 {}, domain_adaptation {}, lambda_dat {}, label_type {}'
.format(batch_size, lr, z_dim, K, dropout, adj_type, learn_edge, lambda1,
lambda2, domain_adaptation, lambda_dat, label_type))
# parameter sanity check
if label_type not in label_types:
raise Exception("undefined label_type")
if adj_type not in adj_types:
raise Exception("undefined adj_type")
# construct model
edge_weight = initial_adjacency_matrix(adj_type)
model = SymSimGCNNet(n_channels, learn_edge, edge_weight, n_bands, [z_dim], n_classes[label_type],
K, dropout, domain_adaptation)
last_epoch = 0
if ckpt_load is not None:
ckpt = torch.load(ckpt_load)
last_epoch = ckpt_load["epoch"]
if last_epoch >= n_epochs:
raise Exception("loaded model have trained >= n_epochs")
state_dict = ckpt_load["state_dict"]
model.load_state_dict(state_dict)
# use multiple GPU
model = DataParallel(model, device_ids=device_ids).to(device)
logger.info(model)
# prepare dataloader
logger.info("tr_dataset: {}".format(tr_dataset))
logger.info("te_dataset: {}".format(te_dataset))
logger.info("training start from epoch {}".format(last_epoch))
tr_loader = DataListLoader(tr_dataset, batch_size, True)
# prepare optimizer
param_list1 = []
param_list2 = []
for name, param in model.named_parameters():
if name in ['module.edge_weight', 'module.conv1.lin.bias', 'module.fc.bias']:
param_list1.append(param)
else:
param_list2.append(param)
optimizer = torch.optim.Adam([
{'params': param_list1, 'weight_decay': 0},
{'params': param_list2, 'weight_decay': lambda2}
], lr=lr)
# iterate over all epochs
eval_acc_list = []
macro_f1_list = []
for ep in range(last_epoch + 1, n_epochs + 1):
model.train()
loss_all = 0
reverse_scale = 2 / (1 + math.exp(-10 * ep / n_epochs)) - 1
if domain_adaptation == 'RevGrad':
model.module.alpha = reverse_scale
# iterate over all graphs
for tr_data_list in tr_loader:
# output shape (len(tr_data_list), 5 or 1)
output, domain_output = model(tr_data_list)
# classification loss
# y shape (len(tr_data_list), )
y = torch.cat([data.y for data in tr_data_list]).to(output.device)
if label_type == "hard":
loss = F.cross_entropy(output, y)
elif label_type == "soft":
loss = - distribution_label(y) * F.log_softmax(output, dim=1)
loss = torch.mean(torch.sum(loss, dim=1))
else:
loss = F.mse_loss(output, y - 2)
# l1 regularization loss
if learn_edge:
loss += lambda1 * torch.sum(torch.abs(model.module.edge_weight))
# domain adaptation loss
if domain_adaptation:
# tr_data.x: [num_graph * n_channels, feature_dim]
n_nodes = domain_output.size(0)
loss += lambda_dat * F.cross_entropy(domain_output, torch.zeros(n_nodes).cuda())
te_indices = torch.randint(0, len(te_dataset), len(tr_data_list))
te_data = te_dataset[te_indices]
_, te_domain_output = model(te_data)
loss += lambda_dat * F.cross_entropy(te_domain_output, torch.ones(n_nodes).cuda())
loss_all += loss.item() * len(tr_data_list)
# optimize the model
optimizer.zero_grad()
loss.backward()
optimizer.step()
# evaluate the model
accuracy, macro_f1_score = evaluate_RGNN(model, te_dataset, label_type)
eval_acc_list.append(accuracy)
macro_f1_list.append(macro_f1_score)
train_acc, _ = evaluate_RGNN(model, tr_dataset, label_type)
logger.info('epoch: {:4d}; loss: {:9.5f}; train acc: {:9.5f}; eval acc: {:9.5f}; '
'macro f1: {:9.5f};'
.format(ep, loss_all/len(tr_dataset), train_acc, accuracy, macro_f1_score))
# save model checkpoint
logger.info(list(model.parameters()))
logger.info(format_list(model.module.edge_weight.detach().cpu().numpy().flatten()))
if ckpt_save_name is not None:
checkpoint = {"epoch": n_epochs, "state_dict": model.state_dict()}
torch.save(checkpoint, ckpt_dir + '/' + ckpt_save_name)
return eval_acc_list, macro_f1_list
def test_model(model_path,
dataset_param, model_param, test_dataset_param,
batch_size=32, layer_num=None,
metric_name_list=['relative_loss', 'mse_loss', 'accuracy',
'post_relative_loss', 'post_mse_loss', 'post_accuracy',
'layer_num',]):
print(model_path)
print(dataset_param, model_param, test_dataset_param)
model_filename = osp.basename(model_path)
new_log_dir = osp.dirname(model_path)
log_dir_base = osp.basename(new_log_dir)
log_dir = osp.dirname(osp.dirname(new_log_dir))
performance_dir = osp.join(log_dir, 'performance')
if not osp.exists(performance_dir):
os.mkdir(performance_dir)
performance_dir = dataset_param2path(performance_dir, test_dataset_param)
new_performance_dir = osp.join(performance_dir, log_dir_base)
if not osp.exists(new_performance_dir):
os.mkdir(new_performance_dir)
batch_size = batch_size
parallel_flag = False
if torch.cuda.is_available():
device = torch.device('cuda')
# parallel_flag = torch.cuda.device_count() > 1
else:
device = torch.device('cpu')
cur_model_path = os.path.join(new_log_dir, model_filename)
cur_performance_dir = osp.join(new_performance_dir, model_filename)
if not osp.exists(cur_performance_dir):
os.mkdir(cur_performance_dir)
if not all([osp.exists(osp.join(cur_performance_dir, ('raw_%s'%metric_name)+('' if layer_num is None else '_'+str(layer_num))+'.csv'))
for metric_name in metric_name_list]):
with torch.no_grad():
test_dataset = param2dataset(test_dataset_param, train_flag=False)
data_loader_fn = DataListLoader if parallel_flag else DataLoader
test_data_loader = data_loader_fn(test_dataset, batch_size)
net = param2model(test_dataset, model_param)
net = net.to(device)
checkpoint = torch.load(cur_model_path, map_location=device)
net.load_state_dict(checkpoint['model_state_dict'])
if layer_num is not None:
try:
net.gnn_module.layer_num = layer_num
except AttributeError:
for i in range(len(net.gnn_module_list)):
net.gnn_module_list[i].layer_num = layer_num
if parallel_flag:
net = DataParallel(net)
test_metric_list, y_list = evaluate(net, test_data_loader, device,
parallel_flag=parallel_flag,
# post_processing_flag=False,
metric_name_list=metric_name_list)
for metric_name, metric_list in test_metric_list.items():
record = np.array([y_list, metric_list])
np.savetxt(osp.join(cur_performance_dir, ('raw_%s'%metric_name)+('' if layer_num is None else '_'+str(layer_num))+'.csv'), record)
del test_dataset, test_data_loader
# gamma: original image ratio
for sigma, gamma in tqdm(
product(
(0.001, 0.01, 0.05, 0.1, 0.3, 0.5, 1, 2, 5, 10),
(0, 0.1, 0.3, 0.5, 0.75, 0.9, 0.99),
)
):
print(
colorama.Fore.MAGENTA
+ "Testing config - sigma: %.3E, gamma: %.3E" % (sigma, gamma)
)
model = NaiveBilateralFilter(fin=6, sigma=sigma, gamma=gamma, k=32)
if parallel:
model = DataParallel(model, device_ids=gpu_ids, output_device=gpu_id).to(
device
)
else:
model = model.to(device)
model.eval()
total_mse, total_psnr, orig_psnr = evaluate(model, train_loader, 0)
records[sigma][gamma] = (total_mse, total_psnr)
if total_psnr > max_psnr:
best_sigma, best_gamma = sigma, gamma
max_psnr, min_mse = total_psnr, total_mse
print(
colorama.Fore.GREEN
+ "Max PSNR: %.3f, min MSE: %.3f, ORIG-MSE: %.3f@ sigma: %.3f, gamma: %.3f"
def test_data_parallel():
module = DataParallel(None)
data_list = [Data(x=torch.randn(x, 1)) for x in [2, 3, 10, 4]]
batches = module.scatter(data_list, device_ids=[0, 1, 0, 1])
assert len(batches) == 3
DataLoader(datasets[2][test_size:test_size + train_size],
batch_size=graphsPerBatch)
]
#print('Can we use GPU? ',torch.cuda.is_available())
# Select which GPUs we can see
#os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#print(device.current_device())
myGCN = multiViewGCN(2, 4, device)
# Use multiple GPUs if we can
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
myGCN = DataParallel(myGCN)
optimizer = torch.optim.Adam(myGCN.parameters(),
lr=0.0005) #, weight_decay=5e-4)
nEpochs = 20
def train():
myGCN.train()
loss_all = 0
loss_func = torch.nn.CrossEntropyLoss()
for data0, data1, data2 in zip(train_loader[0], train_loader[1],
train_loader[2]):
data0 = data0.to(device)
def train():
# set the input channel dims based on featurization type
if args.feature_type == "pybel":
feature_size = 20
else:
feature_size = 75
print("found {} datasets in input train-data".format(len(args.train_data)))
train_dataset_list = []
val_dataset_list = []
for data in args.train_data:
train_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,
use_docking=args.use_docking,
))
for data in args.val_data:
val_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,
use_docking=args.use_docking,
))
train_dataset = ConcatDataset(train_dataset_list)
val_dataset = ConcatDataset(val_dataset_list)
train_dataloader = DataListLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=False,
worker_init_fn=worker_init_fn,
drop_last=True,
) # just to keep batch sizes even, since shuffling is used
val_dataloader = DataListLoader(
val_dataset,
batch_size=args.batch_size,
shuffle=False,
worker_init_fn=worker_init_fn,
drop_last=True,
)
tqdm.write("{} complexes in training dataset".format(len(train_dataset)))
tqdm.write("{} complexes in validation dataset".format(len(val_dataset)))
model = GeometricDataParallel(
PotentialNetParallel(
in_channels=feature_size,
out_channels=1,
covalent_gather_width=args.covalent_gather_width,
non_covalent_gather_width=args.non_covalent_gather_width,
covalent_k=args.covalent_k,
non_covalent_k=args.non_covalent_k,
covalent_neighbor_threshold=args.covalent_threshold,
non_covalent_neighbor_threshold=args.non_covalent_threshold,
)).float()
model.train()
model.to(0)
tqdm.write(str(model))
tqdm.write("{} trainable parameters.".format(
sum(p.numel() for p in model.parameters() if p.requires_grad)))
tqdm.write("{} total parameters.".format(
sum(p.numel() for p in model.parameters())))
criterion = nn.MSELoss().float()
optimizer = Adam(model.parameters(), lr=args.lr)
best_checkpoint_dict = None
best_checkpoint_epoch = 0
best_checkpoint_step = 0
best_checkpoint_r2 = -9e9
step = 0
for epoch in range(args.epochs):
losses = []
for batch in tqdm(train_dataloader):
batch = [x for x in batch if x is not None]
if len(batch) < 1:
print("empty batch, skipping to next batch")
continue
optimizer.zero_grad()
data = [x[2] for x in batch]
y_ = model(data)
y = torch.cat([x[2].y for x in batch])
loss = criterion(y.float(), y_.cpu().float())
losses.append(loss.cpu().data.item())
loss.backward()
y_true = y.cpu().data.numpy()
y_pred = y_.cpu().data.numpy()
r2 = r2_score(y_true=y_true, y_pred=y_pred)
mae = mean_absolute_error(y_true=y_true, y_pred=y_pred)
pearsonr = stats.pearsonr(y_true.reshape(-1), y_pred.reshape(-1))
spearmanr = stats.spearmanr(y_true.reshape(-1), y_pred.reshape(-1))
tqdm.write(
"epoch: {}\tloss:{:0.4f}\tr2: {:0.4f}\t pearsonr: {:0.4f}\tspearmanr: {:0.4f}\tmae: {:0.4f}\tpred stdev: {:0.4f}"
"\t pred mean: {:0.4f} \tcovalent_threshold: {:0.4f} \tnon covalent threshold: {:0.4f}"
.format(
epoch,
loss.cpu().data.numpy(),
r2,
float(pearsonr[0]),
float(spearmanr[0]),
float(mae),
np.std(y_pred),
np.mean(y_pred),
model.module.covalent_neighbor_threshold.t.cpu().data.item(
),
model.module.non_covalent_neighbor_threshold.t.cpu().data.
item(),
))
if args.checkpoint:
if step % args.checkpoint_iter == 0:
checkpoint_dict = checkpoint_model(
model,
val_dataloader,
epoch,
step,
args.checkpoint_dir +
"/model-epoch-{}-step-{}.pth".format(epoch, step),
)
if checkpoint_dict["validate_dict"][
"r2"] > best_checkpoint_r2:
best_checkpoint_step = step
best_checkpoint_epoch = epoch
best_checkpoint_r2 = checkpoint_dict["validate_dict"][
"r2"]
best_checkpoint_dict = checkpoint_dict
optimizer.step()
step += 1
if args.checkpoint:
checkpoint_dict = checkpoint_model(
model,
val_dataloader,
epoch,
step,
args.checkpoint_dir +
"/model-epoch-{}-step-{}.pth".format(epoch, step),
)
if checkpoint_dict["validate_dict"]["r2"] > best_checkpoint_r2:
best_checkpoint_step = step
best_checkpoint_epoch = epoch
best_checkpoint_r2 = checkpoint_dict["validate_dict"]["r2"]
best_checkpoint_dict = checkpoint_dict
if args.checkpoint:
# once broken out of the loop, save last model
checkpoint_dict = checkpoint_model(
model,
val_dataloader,
epoch,
step,
args.checkpoint_dir +
"/model-epoch-{}-step-{}.pth".format(epoch, step),
)
if checkpoint_dict["validate_dict"]["r2"] > best_checkpoint_r2:
best_checkpoint_step = step
best_checkpoint_epoch = epoch
best_checkpoint_r2 = checkpoint_dict["validate_dict"]["r2"]
best_checkpoint_dict = checkpoint_dict
if args.checkpoint:
torch.save(best_checkpoint_dict,
args.checkpoint_dir + "/best_checkpoint.pth")
print("best training checkpoint epoch {}/step {} with r2: {}".format(
best_checkpoint_epoch, best_checkpoint_step, best_checkpoint_r2))
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 !")
args = parser.parse_args()
# Load dataset
dataset = PygNodePropPredDataset(name=args.dataset_name)
data = dataset[0]
dataset_test(data)
if args.multi_gpu:
# Unit test: GPU number verification
# Prepare model
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = parse_model_name(args.model, dataset)
model = DataParallel(model)
model = model.to(device)
#Split graph into subgraphs
if args.subgraph_scheme == 'cluster':
# Split data into subgraphs using cluster methods
data_list = list(ClusterData(data, num_parts=args.num_parts))
elif args.subgraph_scheme == 'neighbor':
data_list = list(
NeighborSubgraphLoader(data,
batch_size=args.neighbor_batch_size))
print(
f'Using neighbor sampling | number of subgraphs: {len(data_list)}'
)
# Run the model for each batch size setups
def model_training(data_list_train, data_list_test, epochs, acc_epoch, acc_epoch2, save_model_epochs, validation_epoch, batchsize, logfilename, load_checkpoint= None):
#logging
logging.basicConfig(level=logging.DEBUG, filename='./logfiles/'+logfilename, filemode="w+",
format="%(message)s")
trainloader = DataListLoader(data_list_train, batch_size=batchsize, shuffle=True)
testloader = DataListLoader(data_list_test, batch_size=batchsize, shuffle=True)
device = torch.device('cuda')
complete_net = completeNet()
complete_net = DataParallel(complete_net)
complete_net = complete_net.to(device)
#train parameters
weights = [10, 1]
optimizer = torch.optim.Adam(complete_net.parameters(), lr=0.001, weight_decay=0.001)
#resume training
initial_epoch=1
if load_checkpoint!=None:
checkpoint = torch.load(load_checkpoint)
complete_net.load_state_dict(checkpoint['model_state_dict'], strict=False)
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
initial_epoch = checkpoint['epoch']+1
loss = checkpoint['loss']
complete_net.train()
for epoch in range(initial_epoch, epochs+1):
epoch_total=0
epoch_total_ones= 0
epoch_total_zeros= 0
epoch_correct=0
epoch_correct_ones= 0
epoch_correct_zeros= 0
running_loss= 0
batches_num=0
for batch in trainloader:
batch_total=0
batch_total_ones= 0
batch_total_zeros= 0
batch_correct= 0
batch_correct_ones= 0
batch_correct_zeros= 0
batches_num+=1
# Forward-Backpropagation
output, output2, ground_truth, ground_truth2, det_num, tracklet_num= complete_net(batch)
optimizer.zero_grad()
loss = weighted_binary_cross_entropy(output, ground_truth, weights)
loss.backward()
optimizer.step()
##Accuracy
if epoch%acc_epoch==0 and epoch!=0:
# Hungarian method, clean up
cleaned_output= hungarian(output2, ground_truth2, det_num, tracklet_num)
batch_total += cleaned_output.size(0)
ones= torch.tensor([1 for x in cleaned_output]).to(device)
zeros = torch.tensor([0 for x in cleaned_output]).to(device)
batch_total_ones += (cleaned_output == ones).sum().item()
batch_total_zeros += (cleaned_output == zeros).sum().item()
batch_correct += (cleaned_output == ground_truth2).sum().item()
temp1 = (cleaned_output == ground_truth2)
temp2 = (cleaned_output == ones)
batch_correct_ones += (temp1 & temp2).sum().item()
temp3 = (cleaned_output == zeros)
batch_correct_zeros += (temp1 & temp3).sum().item()
epoch_total += batch_total
epoch_total_ones += batch_total_ones
epoch_total_zeros += batch_total_zeros
epoch_correct += batch_correct
epoch_correct_ones += batch_correct_ones
epoch_correct_zeros += batch_correct_zeros
if loss.item()!=loss.item():
print("Error")
break
if batch_total_ones != 0 and batch_total_zeros != 0 and epoch%acc_epoch==0 and epoch!=0:
print('Epoch: [%d] | Batch: [%d] | Training_Loss: %.3f | Total_Accuracy: %.3f | Ones_Accuracy: %.3f | Zeros_Accuracy: %.3f |' %
(epoch, batches_num, loss.item(), 100 * batch_correct / batch_total, 100 * batch_correct_ones / batch_total_ones,
100 * batch_correct_zeros / batch_total_zeros))
logging.info('Epoch: [%d] | Batch: [%d] | Training_Loss: %.3f | Total_Accuracy: %.3f | Ones_Accuracy: %.3f | Zeros_Accuracy: %.3f |' %
(epoch, batches_num, loss.item(), 100 * batch_correct / batch_total, 100 * batch_correct_ones / batch_total_ones,
100 * batch_correct_zeros / batch_total_zeros))
else:
print('Epoch: [%d] | Batch: [%d] | Training_Loss: %.3f |' %
(epoch, batches_num, loss.item()))
logging.info('Epoch: [%d] | Batch: [%d] | Training_Loss: %.3f |' %
(epoch, batches_num, loss.item()))
running_loss += loss.item()
if loss.item()!=loss.item():
print("Error")
break
if epoch_total_ones!=0 and epoch_total_zeros!=0 and epoch%acc_epoch==0 and epoch!=0:
print('Epoch: [%d] | Training_Loss: %.3f | Total_Accuracy: %.3f | Ones_Accuracy: %.3f | Zeros_Accuracy: %.3f |' %
(epoch, running_loss / batches_num, 100 * epoch_correct / epoch_total, 100 * \
epoch_correct_ones / epoch_total_ones, 100 * epoch_correct_zeros / epoch_total_zeros))
logging.info('Epoch: [%d] | Training_Loss: %.3f | Total_Accuracy: %.3f | Ones_Accuracy: %.3f | Zeros_Accuracy: %.3f |' %
(epoch, running_loss / batches_num, 100 * epoch_correct / epoch_total, 100 * \
epoch_correct_ones / epoch_total_ones, 100 * epoch_correct_zeros / epoch_total_zeros))
else:
print('Epoch: [%d] | Training_Loss: %.3f |' %
(epoch, running_loss / batches_num))
logging.info('Epoch: [%d] | Training_Loss: %.3f |' %
(epoch, running_loss / batches_num))
# save model
if epoch%save_model_epochs==0 and epoch!=0:
torch.save({
'epoch': epoch,
'model_state_dict': complete_net.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': running_loss,
}, './models/epoch_'+str(epoch)+'.pth')
#validation
if epoch%validation_epoch==0 and epoch!=0:
with torch.no_grad():
epoch_total=0
epoch_total_ones= 0
epoch_total_zeros= 0
epoch_correct=0
epoch_correct_ones= 0
epoch_correct_zeros= 0
running_loss= 0
batches_num=0
for batch in testloader:
batch_total=0
batch_total_ones= 0
batch_total_zeros= 0
batch_correct= 0
batch_correct_ones= 0
batch_correct_zeros= 0
batches_num+=1
output, output2, ground_truth, ground_truth2, det_num, tracklet_num = complete_net(batch)
loss = weighted_binary_cross_entropy(output, ground_truth, weights)
running_loss += loss.item()
##Accuracy
if epoch%acc_epoch2==0 and epoch!=0:
# Hungarian method, clean up
cleaned_output= hungarian(output2, ground_truth2, det_num, tracklet_num)
batch_total += cleaned_output.size(0)
ones= torch.tensor([1 for x in cleaned_output]).to(device)
zeros = torch.tensor([0 for x in cleaned_output]).to(device)
batch_total_ones += (cleaned_output == ones).sum().item()
batch_total_zeros += (cleaned_output == zeros).sum().item()
batch_correct += (cleaned_output == ground_truth2).sum().item()
temp1 = (cleaned_output == ground_truth2)
temp2 = (cleaned_output == ones)
batch_correct_ones += (temp1 & temp2).sum().item()
temp3 = (cleaned_output == zeros)
batch_correct_zeros += (temp1 & temp3).sum().item()
epoch_total += batch_total
epoch_total_ones += batch_total_ones
epoch_total_zeros += batch_total_zeros
epoch_correct += batch_correct
epoch_correct_ones += batch_correct_ones
epoch_correct_zeros += batch_correct_zeros
if epoch_total_ones!=0 and epoch_total_zeros!=0 and epoch%acc_epoch2==0 and epoch!=0:
print('Epoch: [%d] | Validation_Loss: %.3f | Total_Accuracy: %.3f | Ones_Accuracy: %.3f | Zeros_Accuracy: %.3f |' %
(epoch, running_loss / batches_num, 100 * epoch_correct / epoch_total, 100 * \
epoch_correct_ones / epoch_total_ones, 100 * epoch_correct_zeros / epoch_total_zeros))
logging.info('Epoch: [%d] | Validation_Loss: %.3f | Total_Accuracy: %.3f | Ones_Accuracy: %.3f | Zeros_Accuracy: %.3f |' %
(epoch, running_loss / batches_num, 100 * epoch_correct / epoch_total, 100 * \
epoch_correct_ones / epoch_total_ones, 100 * epoch_correct_zeros / epoch_total_zeros))
else:
print('Epoch: [%d] | Validation_Loss: %.3f |' %
(epoch, running_loss / batches_num))
logging.info('Epoch: [%d] | Validation_Loss: %.3f |' %
(epoch, running_loss / batches_num))
def GCN(dataset, params, num_pre_epochs, num_epochs, MonteSize, PruningTimes,
width, lr, savepath):
Batch_size = int(params[0])
for Monte_iter in range(MonteSize):
# Data
best_loss = float('inf') # best test loss
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
TrainConvergence = []
TestConvergence = []
# model
root = '/git/data/GraphData/' + dataset
if dataset == 'Cora':
model_name = "PruningGCN"
datasetroot = Planetoid(root=root, name=dataset).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-Monte_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], params[2],
params[3], Monte_iter)
if Monte_iter == 0:
if resume == True and os.path.exists(model_to_save):
[
OptimizedNet, NewNetworksize, TrainConvergence,
TestConvergence, start_epoch
] = ResumeModel(model_to_save)
if start_epoch >= num_epochs - 1:
continue
else:
net = Net(datasetroot, width)
#net.apply(weight_reset)
elif dataset == 'ENZYMES' or dataset == 'MUTAG':
model_name = "topk_pool_Net"
datasetroot = TUDataset(root, name=dataset)
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], params[2],
params[3])
if Monte_iter == 0:
if resume == True and os.path.exists(model_to_save):
[
OptimizedNet, NewNetworksize, TrainConvergence,
TestConvergence, start_epoch
] = ResumeModel(model_to_save)
if start_epoch >= num_epochs - 1:
continue
else:
net = topk_pool_Net(datasetroot, width)
elif dataset == 'MNIST':
datasetroot = MNISTSuperpixels(root=root,
transform=T.Cartesian()).shuffle()
trainloader = DataListLoader(datasetroot,
batch_size=Batch_size,
shuffle=True)
testloader = DataListLoader(datasetroot,
batch_size=100,
shuffle=False)
model_name = 'SPlineNet'
model_to_save = './checkpoint/{}-{}-param_{}_{}_{}_{}-Mon_{}-ckpt.pth'.format(
dataset, model_name, params[0], params[1], params[2],
params[3], Monte_iter)
if resume == "True" and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= num_epochs - 1:
continue
else:
#net=Net(datasetroot,width)
net = SPlineNet(datasetroot, width)
elif dataset == 'CIFAR10':
if resume == "True" and os.path.exists(model_to_save):
[net, TrainConvergence, TestConvergence,
start_epoch] = ResumeModel(model_to_save)
if start_epoch >= num_epochs - 1:
continue
else:
net = getattr(CIFAR10_resnet, 'Resnet20_CIFAR10')(params[1])
else:
raise Exception(
"The dataset is:{}, it isn't existed.".format(dataset))
if Monte_iter == 0 and start_epoch == 0:
print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
net = DataParallel(net)
device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
net = net.to(device)
#cudnn.benchmark = True
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(),
lr=lr,
momentum=0.9,
weight_decay=5e-4)
logging(
'Batch size: {}, Number of layers:{} ConCoeff: {}, CutoffCoffi:{}, MonteSize:{}'
.format(params[0], params[1], params[2], params[3],
Monte_iter))
for epoch in range(num_pre_epochs):
PreTrainLoss = train(trainloader, net, optimizer, criterion)
print('\nEpoch: {}, Average pre-tain loss: {:.4f} \n'.format(
epoch, PreTrainLoss[0]))
NewNetworksize = RetainNetworkSize(net, params[2])
del net
#NewNetworksize=width
for pruningIter in range(PruningTimes):
if pruningIter > 0:
[
OptimizedNet, NewNetworksize, TrainConvergence,
TestConvergence, start_epoch
] = ResumeModel(model_to_save)
elif dataset == 'Cora' and start_epoch == 0:
OptimizedNet = Net(datasetroot, NewNetworksize[0:-1])
elif dataset == 'ENZYMES' and start_epoch == 0:
NewNetworkSizeAdjust = NewNetworksize[0:-1]
NewNetworkSizeAdjust[0] = width[0] - 1
OptimizedNet = topk_pool_Net(datasetroot, NewNetworkSizeAdjust)
OptimizedNet = DataParallel(OptimizedNet)
device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
OptimizedNet = OptimizedNet.to(device)
cudnn.benchmark = True
criterionNew = nn.CrossEntropyLoss()
optimizerNew = optim.SGD(OptimizedNet.parameters(),
lr=lr,
momentum=0.9,
weight_decay=5e-4)
criterion = nn.CrossEntropyLoss()
for epoch in range(start_epoch, num_epochs):
TrainLoss = train(trainloader, OptimizedNet, optimizerNew,
criterionNew)
print('\n Epoch: {}, Average tain loss: {:.4f} \n'.format(
epoch, TrainLoss[0]))
TrainConvergence.append(statistics.mean(TrainLoss))
NewNetworksize = RetainNetworkSize(OptimizedNet, params[2])
TestConvergence.append(
statistics.mean(test(testloader, OptimizedNet, criterion)))
# save model
if TestConvergence[epoch] < best_loss:
logging('Saving..')
state = {
'net': OptimizedNet.module,
'TrainConvergence': TrainConvergence,
'TestConvergence': TestConvergence,
'epoch': num_epochs,
'NewNetworksize': NewNetworksize[0:-1],
}
if not os.path.isdir('checkpoint'):
os.mkdir('checkpoint')
torch.save(state, model_to_save)
best_loss = TestConvergence[epoch]
## save recurrence plots
"""if epoch%20==0:
save_recurrencePlots_file="../Results/RecurrencePlots/RecurrencePlots_{}_{}_BatchSize{} \_ConCoeffi{}_epoch{}.png".format(dataset, model_name,params[0],params[1],epoch)
save_recurrencePlots(net,save_recurrencePlots_file)"""
FileName = "{}-{}-param_{}_{}_{}_{}-monte_{}".format(
dataset, model_name, params[0], params[1], params[2], params[3],
Monte_iter)
np.save(savepath + 'TrainConvergence-' + FileName, TrainConvergence)
#np.save(savepath+'TestConvergence-'+FileName,TestConvergence)
#torch.cuda.empty_cache()
print_nvidia_useage()
if return_output == True:
return TestConvergence[-1], net.module.fc.weight
else:
pass
