def main():
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
print("Load dataset...")
DATA_PATH = config.DATA['root']
TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH,
npoint=config.SETTING['num_point'],
split='train',
normal_channel=config.SETTING['normal'])
TEST_DATASET = ModelNetDataLoader(root=DATA_PATH,
npoint=config.SETTING['num_point'],
split='test',
normal_channel=config.SETTING['normal'])
trainDataLoader = torch.utils.data.DataLoader(
TRAIN_DATASET,
batch_size=config.TRAIN['batch_size'],
shuffle=True,
num_workers=4)
testDataLoader = torch.utils.data.DataLoader(
TEST_DATASET,
batch_size=config.TRAIN['batch_size'],
shuffle=False,
num_workers=4)
print("ok!")
print("Check device...")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
print("ok!")
print("Load model...")
if config.TRAIN['model'] == "POINTNET":
model = PointNet(config.DATA['num_class'],
normal_channel=config.SETTING['normal'])
elif config.TRAIN['model'] == "DGCNN":
model = DGCNN(config.DATA['num_class'])
model = model.to(device)
print("ok!")
print("Load optimizer...")
optimizer = torch.optim.Adam(model.parameters(),
lr=config.TRAIN['learning_rate'],
betas=(0.9, 0.999),
eps=1e-08,
weight_decay=config.TRAIN['decay_rate'])
print("ok!")
# Starting training
print("Starting training... ")
best = 1
for epoch in range(config.TRAIN['epoch']):
train(epoch, trainDataLoader, optimizer, model, device)
acc = valid(testDataLoader, model, device)
pred = torch.argmax(pred, dim=-1)
pred, y = pred.cpu().detach().numpy(), batch.y.detach().numpy()
correct.append(np.mean(pred == y))
return np.mean(correct)
DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device(
'cpu')
pre_transform, transform = T.NormalizeScale(), T.SamplePoints(1024)
dataset = ModelNet('./data/model-net', '10', True, transform, pre_transform)
dataloader = DataLoader(dataset, num_workers=2, batch_size=16, shuffle=True)
loss_dict = {}
for lr in (5e-3, ):
model = PointNet(3, [64, 64], [128, 128, 128], 10).to(DEVICE)
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
key = lr
loss_dict[key] = []
for epoch in range(5):
loss = train_epoch()
loss_dict[key].append(loss)
acc = test()
print(f'Key {key} epoch {epoch}, loss {loss: 0.2f}, acc {acc: 0.2f}')
print()
plt.figure()
for k in loss_dict.keys():
def main():
#Parse Arguments
parser = argparse.ArgumentParser(
description=
'Uses MDAnalysis and PointNet to identify largest cluster of solid-like atoms'
)
parser.add_argument('--weights',
help='folder containing network weights to use',
type=str,
required=True)
parser.add_argument('--nclass',
help='number of classes',
type=int,
required=True)
parser.add_argument('--trjpath',
help='path to gro/xtc files',
type=str,
required=True)
parser.add_argument('--cutoff',
help='neighbor cutoff distance (in nm)',
type=float,
required=True)
parser.add_argument('--maxneigh',
help='max number of neighbors',
type=int,
required=True)
parser.add_argument('--outname',
help='output file name',
type=str,
required=True)
args = parser.parse_args()
# Import topology
u = mda.Universe(args.trjpath + '.gro', args.trjpath + '.xtc')
# File to write output
f_summary = open(args.outname + '_summary.mda', 'w')
f_class = open(args.outname + '_class.mda', 'w')
f_summary.write(
"# Time, Largest_cls_solid, n_liq, n_fcc, n_hcp, n_bcc, n_solid\n")
# Here is where we initialize the pointnet
pointnet = PointNet(n_points=args.maxneigh,
n_classes=args.nclass,
weights_dir=args.weights)
# Loop over trajectory
for ts in u.trajectory:
# Generate neighbor list
print("Generating neighbor list")
nlist = nsgrid.FastNS(args.cutoff * 10.0, u.atoms.positions,
ts.dimensions).self_search()
# Extract required info
ndxs = nlist.get_indices()
dxs = nlist.get_dx()
dists = nlist.get_distances()
print("Extracted all relevant information")
samples = []
# Prepare samples to send through pointnet
for i in range(len(dxs)):
## Sort neighbors by distance (so that we can
## normalize all distances such that the distance
## to the closest neighbor is 1.0 )
nneigh = int(len(dxs[i]) / 3)
np_dxs = np.asarray(dxs[i]).reshape([nneigh, 3])
sort_order = np.asarray(dists[i]).argsort()
# Sort neighbors by distance
np_dxs = np_dxs[sort_order]
if nneigh > 0:
np_dxs /= np.linalg.norm(np_dxs[0])
# Now correctly size/pad the point cloud
if nneigh < args.maxneigh:
np_dxs = np.pad(
np_dxs,
[(0, args.maxneigh - nneigh), (0, 0)],
'constant',
)
elif nneigh > args.maxneigh:
np_dxs = np_dxs[:args.maxneigh]
# Append sample info
samples.append(np_dxs)
# And convert to np array
np_samples = np.asarray(samples)
print("Frame {}, Shape sent to pointnet: {}".format(
ts.frame, np_samples.shape))
sys.stdout.flush()
# Send sample through inference
results = pointnet.infer_nolabel(np_samples)
results = np.asarray(results)
print("Frame {}, Results returned from pointnet, shape {}".format(
ts.frame, results.shape))
sys.stdout.flush()
# Extract different atom types
liquid_atoms = np.where(results == 0)[0]
fcc_atoms = np.where(results == 1)[0]
hcp_atoms = np.where(results == 2)[0]
bcc_atoms = np.where(results == 3)[0]
solid_atoms = np.where(results > 0)[0]
print("%d total solid atoms" % solid_atoms.shape[0])
## Now we are going to construct the largest cluster of
## solid atoms in the system (i.e., a solid nucleus)
# We need neighbor lists for connectivity cutoff
# Using 15.0 Angstroms (mda units) here
nlist = nsgrid.FastNS(15.0, u.atoms.positions,
ts.dimensions).self_search()
pairs = nlist.get_pairs()
# Find the largest cluster of solids (not liquid)
G = nx.Graph()
G.add_edges_from(pairs)
G.remove_nodes_from(liquid_atoms)
largest_cluster = max(nx.connected_component_subgraphs(G), key=len)
f_summary.write("{:8.3f}{:8d}{:8d}{:8d}{:8d}{:8d}{:8d}\n".format(
ts.time, len(largest_cluster), liquid_atoms.shape[0],
fcc_atoms.shape[0], hcp_atoms.shape[0], bcc_atoms.shape[0],
solid_atoms.shape[0]))
for node in largest_cluster:
f_class.write("{:10d}{:8d}{:8d}\n".format(ts.frame + 1, node,
results[node]))
f_summary.close()
f_class.close()
def main(config):
set_seed(config['seed'])
results_dir = prepare_results_dir(config, 'aae', 'training')
starting_epoch = find_latest_epoch(results_dir) + 1
if not exists(join(results_dir, 'config.json')):
with open(join(results_dir, 'config.json'), mode='w') as f:
json.dump(config, f)
setup_logging(results_dir)
log = logging.getLogger('aae')
device = cuda_setup(config['cuda'], config['gpu'])
log.info(f'Device variable: {device}')
if device.type == 'cuda':
log.info(f'Current CUDA device: {torch.cuda.current_device()}')
weights_path = join(results_dir, 'weights')
metrics_path = join(results_dir, 'metrics')
#
# Dataset
#
dataset_name = config['dataset'].lower()
if dataset_name == 'shapenet':
from datasets.shapenet import ShapeNetDataset
dataset = ShapeNetDataset(root_dir=config['data_dir'],
classes=config['classes'])
else:
raise ValueError(f'Invalid dataset name. Expected `shapenet` or '
f'`faust`. Got: `{dataset_name}`')
log.info("Selected {} classes. Loaded {} samples.".format(
'all' if not config['classes'] else ','.join(config['classes']),
len(dataset)))
points_dataloader = DataLoader(dataset, batch_size=config['batch_size'],
shuffle=config['shuffle'],
num_workers=config['num_workers'],
pin_memory=True)
pointnet = config.get('pointnet', False)
#
# Models
#
hyper_network = aae.HyperNetwork(config, device).to(device)
if pointnet:
from models.pointnet import PointNet
encoder = PointNet(config).to(device)
# PointNet initializes it's own weights during instance creation
else:
encoder = aae.Encoder(config).to(device)
encoder.apply(weights_init)
discriminator = aae.Discriminator(config).to(device)
hyper_network.apply(weights_init)
discriminator.apply(weights_init)
if config['reconstruction_loss'].lower() == 'chamfer':
if pointnet:
from utils.metrics import chamfer_distance
reconstruction_loss = chamfer_distance
else:
from losses.champfer_loss import ChamferLoss
reconstruction_loss = ChamferLoss().to(device)
elif config['reconstruction_loss'].lower() == 'earth_mover':
from utils.metrics import earth_mover_distance
reconstruction_loss = earth_mover_distance
else:
raise ValueError(f'Invalid reconstruction loss. Accepted `chamfer` or '
f'`earth_mover`, got: {config["reconstruction_loss"]}')
#
# Optimizers
#
e_hn_optimizer = getattr(optim, config['optimizer']['E_HN']['type'])
e_hn_optimizer = e_hn_optimizer(chain(encoder.parameters(), hyper_network.parameters()),
**config['optimizer']['E_HN']['hyperparams'])
discriminator_optimizer = getattr(optim, config['optimizer']['D']['type'])
discriminator_optimizer = discriminator_optimizer(discriminator.parameters(),
**config['optimizer']['D']['hyperparams'])
log.info("Starting epoch: %s" % starting_epoch)
if starting_epoch > 1:
log.info("Loading weights...")
hyper_network.load_state_dict(torch.load(
join(weights_path, f'{starting_epoch - 1:05}_G.pth')))
encoder.load_state_dict(torch.load(
join(weights_path, f'{starting_epoch - 1:05}_E.pth')))
discriminator.load_state_dict(torch.load(
join(weights_path, f'{starting_epoch-1:05}_D.pth')))
e_hn_optimizer.load_state_dict(torch.load(
join(weights_path, f'{starting_epoch - 1:05}_EGo.pth')))
discriminator_optimizer.load_state_dict(torch.load(
join(weights_path, f'{starting_epoch-1:05}_Do.pth')))
log.info("Loading losses...")
losses_e = np.load(join(metrics_path, f'{starting_epoch - 1:05}_E.npy')).tolist()
losses_g = np.load(join(metrics_path, f'{starting_epoch - 1:05}_G.npy')).tolist()
losses_eg = np.load(join(metrics_path, f'{starting_epoch - 1:05}_EG.npy')).tolist()
losses_d = np.load(join(metrics_path, f'{starting_epoch - 1:05}_D.npy')).tolist()
else:
log.info("First epoch")
losses_e = []
losses_g = []
losses_eg = []
losses_d = []
normalize_points = config['target_network_input']['normalization']['enable']
if normalize_points:
normalization_type = config['target_network_input']['normalization']['type']
assert normalization_type == 'progressive', 'Invalid normalization type'
target_network_input = None
for epoch in range(starting_epoch, config['max_epochs'] + 1):
start_epoch_time = datetime.now()
log.debug("Epoch: %s" % epoch)
hyper_network.train()
encoder.train()
discriminator.train()
total_loss_all = 0.0
total_loss_reconstruction = 0.0
total_loss_encoder = 0.0
total_loss_discriminator = 0.0
total_loss_regularization = 0.0
for i, point_data in enumerate(points_dataloader, 1):
X, _ = point_data
X = X.to(device)
# Change dim [BATCH, N_POINTS, N_DIM] -> [BATCH, N_DIM, N_POINTS]
if X.size(-1) == 3:
X.transpose_(X.dim() - 2, X.dim() - 1)
if pointnet:
_, feature_transform, codes = encoder(X)
else:
codes, _, _ = encoder(X)
# discriminator training
noise = torch.empty(codes.shape[0], config['z_size']).normal_(mean=config['normal_mu'],
std=config['normal_std']).to(device)
synth_logit = discriminator(codes)
real_logit = discriminator(noise)
if config.get('wasserstein', True):
loss_discriminator = torch.mean(synth_logit) - torch.mean(real_logit)
alpha = torch.rand(codes.shape[0], 1).to(device)
differences = codes - noise
interpolates = noise + alpha * differences
disc_interpolates = discriminator(interpolates)
# gradient_penalty_function
gradients = grad(
outputs=disc_interpolates,
inputs=interpolates,
grad_outputs=torch.ones_like(disc_interpolates).to(device),
create_graph=True,
retain_graph=True,
only_inputs=True)[0]
slopes = torch.sqrt(torch.sum(gradients ** 2, dim=1))
gradient_penalty = ((slopes - 1) ** 2).mean()
loss_gp = config['gradient_penalty_coef'] * gradient_penalty
loss_discriminator += loss_gp
else:
# An alternative is a = -1, b = 1 iff c = 0
a = 0.0
b = 1.0
loss_discriminator = 0.5 * ((real_logit - b)**2 + (synth_logit - a)**2)
discriminator_optimizer.zero_grad()
discriminator.zero_grad()
loss_discriminator.backward(retain_graph=True)
total_loss_discriminator += loss_discriminator.item()
discriminator_optimizer.step()
# hyper network training
target_networks_weights = hyper_network(codes)
X_rec = torch.zeros(X.shape).to(device)
for j, target_network_weights in enumerate(target_networks_weights):
target_network = aae.TargetNetwork(config, target_network_weights).to(device)
if not config['target_network_input']['constant'] or target_network_input is None:
target_network_input = generate_points(config=config, epoch=epoch, size=(X.shape[2], X.shape[1]))
X_rec[j] = torch.transpose(target_network(target_network_input.to(device)), 0, 1)
if pointnet:
loss_reconstruction = config['reconstruction_coef'] * \
reconstruction_loss(torch.transpose(X, 1, 2).contiguous(),
torch.transpose(X_rec, 1, 2).contiguous(),
batch_size=X.shape[0]).mean()
else:
loss_reconstruction = torch.mean(
config['reconstruction_coef'] *
reconstruction_loss(X.permute(0, 2, 1) + 0.5,
X_rec.permute(0, 2, 1) + 0.5))
# encoder training
synth_logit = discriminator(codes)
if config.get('wasserstein', True):
loss_encoder = -torch.mean(synth_logit)
else:
# An alternative is c = 0 iff a = -1, b = 1
c = 1.0
loss_encoder = 0.5 * (synth_logit - c)**2
if pointnet:
regularization_loss = config['feature_regularization_coef'] * \
feature_transform_regularization(feature_transform).mean()
loss_all = loss_reconstruction + loss_encoder + regularization_loss
else:
loss_all = loss_reconstruction + loss_encoder
e_hn_optimizer.zero_grad()
encoder.zero_grad()
hyper_network.zero_grad()
loss_all.backward()
e_hn_optimizer.step()
total_loss_reconstruction += loss_reconstruction.item()
total_loss_encoder += loss_encoder.item()
total_loss_all += loss_all.item()
if pointnet:
total_loss_regularization += regularization_loss.item()
log.info(
f'[{epoch}/{config["max_epochs"]}] '
f'Total_Loss: {total_loss_all / i:.4f} '
f'Loss_R: {total_loss_reconstruction / i:.4f} '
f'Loss_E: {total_loss_encoder / i:.4f} '
f'Loss_D: {total_loss_discriminator / i:.4f} '
f'Time: {datetime.now() - start_epoch_time}'
)
if pointnet:
log.info(f'Loss_Regularization: {total_loss_regularization / i:.4f}')
losses_e.append(total_loss_reconstruction)
losses_g.append(total_loss_encoder)
losses_eg.append(total_loss_all)
losses_d.append(total_loss_discriminator)
#
# Save intermediate results
#
if epoch % config['save_samples_frequency'] == 0:
log.debug('Saving samples...')
X = X.cpu().numpy()
X_rec = X_rec.detach().cpu().numpy()
for k in range(min(5, X_rec.shape[0])):
fig = plot_3d_point_cloud(X_rec[k][0], X_rec[k][1], X_rec[k][2], in_u_sphere=True, show=False,
title=str(epoch))
fig.savefig(join(results_dir, 'samples', f'{epoch}_{k}_reconstructed.png'))
plt.close(fig)
fig = plot_3d_point_cloud(X[k][0], X[k][1], X[k][2], in_u_sphere=True, show=False)
fig.savefig(join(results_dir, 'samples', f'{epoch}_{k}_real.png'))
plt.close(fig)
if config['clean_weights_dir']:
log.debug('Cleaning weights path: %s' % weights_path)
shutil.rmtree(weights_path, ignore_errors=True)
os.makedirs(weights_path, exist_ok=True)
if epoch % config['save_weights_frequency'] == 0:
log.debug('Saving weights and losses...')
torch.save(hyper_network.state_dict(), join(weights_path, f'{epoch:05}_G.pth'))
torch.save(encoder.state_dict(), join(weights_path, f'{epoch:05}_E.pth'))
torch.save(e_hn_optimizer.state_dict(), join(weights_path, f'{epoch:05}_EGo.pth'))
torch.save(discriminator.state_dict(), join(weights_path, f'{epoch:05}_D.pth'))
torch.save(discriminator_optimizer.state_dict(), join(weights_path, f'{epoch:05}_Do.pth'))
np.save(join(metrics_path, f'{epoch:05}_E'), np.array(losses_e))
np.save(join(metrics_path, f'{epoch:05}_G'), np.array(losses_g))
np.save(join(metrics_path, f'{epoch:05}_EG'), np.array(losses_eg))
np.save(join(metrics_path, f'{epoch:05}_D'), np.array(losses_d))
batch_size=64,
n_input=dc.train.data.shape[-1],
verbose=1)
if args.net == 'cnn':
net = CNN(epochs=args.nepochs,
batch_size=64,
n_input=dc.train.data.shape[-1],
verbose=1)
if args.net == 'pc':
net = PointNet(epochs=args.nepochs,
batch_size=64,
n_points=dc.train.data.shape[1],
n_classes=dc.train.labels.shape[-1],
n_input=dc.train.data.shape[-1],
verbose=1,
save=1,
noise=args.noise,
params=[args.p1, args.p2],
weights_dir=args.weights)
print('train shape: ' + str(dc.train.data.shape))
print('label shape: ' + str(dc.train.labels.shape))
print('test shape: ' + str(dc.test.data.shape))
print('label shape: ' + str(dc.test.labels.shape))
acc = net.run(dc)
print('final accuracy: ' + str(acc))
def main():
nb_classes = 40
train_file = './ModelNet40/ply_data_train.h5'
test_file = './ModelNet40/ply_data_test.h5'
epochs = 80
fine_tune_epochs = 20
batch_size = 32
train = DataGenerator(train_file, batch_size, nb_classes, train=True)
val = DataGenerator(test_file, batch_size, nb_classes, train=False)
model = PointNet(nb_classes)
model.summary()
lr = 0.001
adam = Adam(lr=lr)
model.compile(optimizer=adam,
loss='categorical_crossentropy',
metrics=['accuracy'])
if not os.path.exists('./results/'):
os.mkdir('./results/')
learning_rate_scheduler = LearningRateScheduler(schedule=schedule)
# pre-train
history = model.fit_generator(train.generator(),
steps_per_epoch=9840 // batch_size,
epochs=epochs,
validation_data=val.generator(),
validation_steps=2468 // batch_size,
callbacks=[learning_rate_scheduler],
verbose=1)
save_history(history, './results/')
model.save_weights('./results/pointnet_weights.h5')
# prune weights
# save masks for weight layers
masks = {}
layer_count = 0
# not compress first convolution layer
first_conv = True
for layer in model.layers:
weight = layer.get_weights()
if len(weight) >= 2:
if not first_conv:
w = deepcopy(weight)
tmp, mask = prune_weights(w[0],
compress_rate=args.compress_rate)
masks[layer_count] = mask
w[0] = tmp
layer.set_weights(w)
else:
first_conv = False
layer_count += 1
# evaluate model after pruning
model.compile(optimizer=Adam(lr=0.0001),
loss='categorical_crossentropy',
metrics=['accuracy'])
score = model.evaluate_generator(val.data_generator(),
steps=2468 // batch_size)
print('val loss: {}'.format(score[0]))
print('val acc: {}'.format(score[1]))
# fine-tune
for i in range(fine_tune_epochs):
for X, Y in train.data_generator():
# train on each batch
model.train_on_batch(X, Y)
# apply masks
for layer_id in masks:
w = model.layers[layer_id].get_weights()
w[0] = w[0] * masks[layer_id]
model.layers[layer_id].set_weights(w)
score = model.evaluate_generator(val.data_generator(),
steps=2468 // batch_size)
print('val loss: {}'.format(score[0]))
print('val acc: {}'.format(score[1]))
# save compressed weights
compressed_name = './results/compressed_pointnet_weights'
save_compressed_weights(model, compressed_name)
parser = argparse.ArgumentParser(
description='Run inference on specified dataset')
parser.add_argument('--weights',
help='folder containing network weights to use',
type=str,
required=True)
parser.add_argument('--dataset',
help='dataset to be used (numpy format)',
type=str,
required=True)
parser.add_argument('--labels',
help='labels corresponding to dataset (numpy format)',
type=str,
required=False)
args = parser.parse_args()
# Load data and labels --
d = np.load(args.dataset)
l = np.load(args.labels)
ic = IC(data=d, labels=l)
pc = PointNet(n_points=ic.test.data.shape[1],
n_classes=ic.test.labels.shape[-1],
n_input=3,
weights_dir=args.weights)
acc = pc.inference(ic, conf_matrix=True)
print('inference accuracy: ' + str(acc))
if not os.path.exists(args.weights):
print('creating weight directory ' + str(args.weights))
os.makedirs(args.weights)
print('loading numpy data...')
data = np.load(args.dataset)
labels = np.load(args.labels)
print('converting to DataContainer format...')
dc = DC(data=data, labels=labels)
net = PointNet(epochs=args.nepochs,
batch_size=64,
lr=args.learning_rate,
n_points=dc.train.data.shape[1],
n_classes=dc.train.labels.shape[-1],
n_input=dc.train.data.shape[-1],
verbose=1,
save=1,
weights_dir=args.weights)
print('train shape: ' + str(dc.train.data.shape))
print('label shape: ' + str(dc.train.labels.shape))
print('test shape: ' + str(dc.test.data.shape))
print('label shape: ' + str(dc.test.labels.shape))
acc = net.run(dc)
print('final accuracy: ' + str(acc))
from torch.utils.data import DataLoader
from torch import nn, optim
from models.pointnet import PointNet
from dataloader import ModelNet40
import torch
BATCH_SIZE = 64
NUM_POINTS = 2048 # 点群の点数
NUM_LABELS = 40 # labelの種類数
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = PointNet(NUM_POINTS, NUM_LABELS)
net.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
loss_l = []
accurary_l = []
trainloader = DataLoader(ModelNet40("data/modelnet40_ply_hdf5_2048/"),
batch_size=64,
shuffle=True)
for epoch in range(2):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data[0].to(device), data[1].to(device)
optimizer.zero_grad()
outputs = net(inputs.view(-1, 3))
parser.add_argument('--nclass', help='number of classes', type=int, required=True)
args = parser.parse_args()
# Load data and labels --
d = np.load(args.dataset)
# Extract just coords
dataset = d[:,:,2:]
print("Dataset shape:")
print(dataset.shape)
nsamples = dataset.shape[0]
npoints = dataset.shape[1]
# Extract just atom/frids
ids = d[:,0,:2]
pc = PointNet(n_points=npoints, n_classes=args.nclass, weights_dir=args.weights)
result = pc.infer_nolabel(dataset)
np_result = np.asarray(result)
np_result = np_result.reshape((nsamples,1))
print("Result shape:")
print(np_result.shape)
print("Ids shape:")
print(ids.shape)
final = np.hstack((ids,np_result))
f = open("classification.out",'w')
for i in range(len(final)):
# Type '0' is liquid -- don't write this to keep classification.out