dmlc--dgl
Mufei Li
134 行
4.4 KiB
Python
134 行
4.4 KiB
Python
import datetime
|
|
import matplotlib.pyplot as plt
|
|
import os
|
|
import random
|
|
import torch
|
|
import torch.backends.cudnn as cudnn
|
|
import torch.nn as nn
|
|
import torch.nn.init as init
|
|
from pprint import pprint
|
|
|
|
|
|
########################################################################################################################
|
|
# configuration #
|
|
########################################################################################################################
|
|
|
|
def mkdir_p(path):
|
|
import errno
|
|
try:
|
|
os.makedirs(path)
|
|
print('Created directory {}'.format(path))
|
|
except OSError as exc:
|
|
if exc.errno == errno.EEXIST and os.path.isdir(path):
|
|
print('Directory {} already exists.'.format(path))
|
|
else:
|
|
raise
|
|
|
|
def date_filename(base_dir='./'):
|
|
dt = datetime.datetime.now()
|
|
return os.path.join(base_dir, '{}_{:02d}-{:02d}-{:02d}'.format(
|
|
dt.date(), dt.hour, dt.minute, dt.second
|
|
))
|
|
|
|
def setup_log_dir(opts):
|
|
log_dir = '{}'.format(date_filename(opts['log_dir']))
|
|
mkdir_p(log_dir)
|
|
return log_dir
|
|
|
|
def save_arg_dict(opts, filename='settings.txt'):
|
|
def _format_value(v):
|
|
if isinstance(v, float):
|
|
return '{:.4f}'.format(v)
|
|
elif isinstance(v, int):
|
|
return '{:d}'.format(v)
|
|
else:
|
|
return '{}'.format(v)
|
|
|
|
save_path = os.path.join(opts['log_dir'], filename)
|
|
with open(save_path, 'w') as f:
|
|
for key, value in opts.items():
|
|
f.write('{}\t{}\n'.format(key, _format_value(value)))
|
|
print('Saved settings to {}'.format(save_path))
|
|
|
|
def setup(args):
|
|
opts = args.__dict__.copy()
|
|
|
|
cudnn.benchmark = False
|
|
cudnn.deterministic = True
|
|
|
|
# Seed
|
|
if opts['seed'] is None:
|
|
opts['seed'] = random.randint(1, 10000)
|
|
random.seed(opts['seed'])
|
|
torch.manual_seed(opts['seed'])
|
|
|
|
# Dataset
|
|
from configure import dataset_based_configure
|
|
opts = dataset_based_configure(opts)
|
|
|
|
assert opts['path_to_dataset'] is not None, 'Expect path to dataset to be set.'
|
|
if not os.path.exists(opts['path_to_dataset']):
|
|
if opts['dataset'] == 'cycles':
|
|
from cycles import generate_dataset
|
|
generate_dataset(opts['min_size'], opts['max_size'],
|
|
opts['ds_size'], opts['path_to_dataset'])
|
|
else:
|
|
raise ValueError('Unsupported dataset: {}'.format(opts['dataset']))
|
|
|
|
# Optimization
|
|
if opts['clip_grad']:
|
|
assert opts['clip_grad'] is not None, 'Expect the gradient norm constraint to be set.'
|
|
|
|
# Log
|
|
print('Prepare logging directory...')
|
|
log_dir = setup_log_dir(opts)
|
|
opts['log_dir'] = log_dir
|
|
mkdir_p(log_dir + '/samples')
|
|
|
|
plt.switch_backend('Agg')
|
|
|
|
save_arg_dict(opts)
|
|
pprint(opts)
|
|
|
|
return opts
|
|
|
|
########################################################################################################################
|
|
# model #
|
|
########################################################################################################################
|
|
|
|
def weights_init(m):
|
|
'''
|
|
Code from https://gist.github.com/jeasinema/ed9236ce743c8efaf30fa2ff732749f5
|
|
Usage:
|
|
model = Model()
|
|
model.apply(weight_init)
|
|
'''
|
|
if isinstance(m, nn.Linear):
|
|
init.xavier_normal_(m.weight.data)
|
|
init.normal_(m.bias.data)
|
|
elif isinstance(m, nn.GRUCell):
|
|
for param in m.parameters():
|
|
if len(param.shape) >= 2:
|
|
init.orthogonal_(param.data)
|
|
else:
|
|
init.normal_(param.data)
|
|
|
|
def dgmg_message_weight_init(m):
|
|
"""
|
|
This is similar as the function above where we initialize linear layers from a normal distribution with std
|
|
1./10 as suggested by the author. This should only be used for the message passing functions, i.e. fe's in the
|
|
paper.
|
|
"""
|
|
def _weight_init(m):
|
|
if isinstance(m, nn.Linear):
|
|
init.normal_(m.weight.data, std=1./10)
|
|
init.normal_(m.bias.data, std=1./10)
|
|
else:
|
|
raise ValueError('Expected the input to be of type nn.Linear!')
|
|
|
|
if isinstance(m, nn.ModuleList):
|
|
for layer in m:
|
|
layer.apply(_weight_init)
|
|
else:
|
|
m.apply(_weight_init)
|