def main():
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
# ----------
# Model and Optimizer
# ----------
model = Net(in_channels=3).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
1000,
gamma=0.5,
last_epoch=-1)
# ----------
# Data
# ----------
pairDataLoader = torch.utils.data.DataLoader(
PairDataset(datasets.ImageFolder(root='./data/raw'),
datasets.ImageFolder(root='./data/contour'),
size=SIZE),
batch_size=BATCH_SIZE,
shuffle=True,
drop_last=True,
pin_memory=True,
)
pairDataLoaderIter = iter(pairDataLoader)
# ----------
# Training
# ----------
for iteration in range(1, ITERATIONS + 1):
try:
rawData, contourData, rawFullData, rawFilteredData = next(
pairDataLoaderIter)
except StopIteration:
pairDataLoaderIter = iter(pairDataLoader)
rawData, contourData, rawFullData, rawFilteredData = next(
pairDataLoaderIter)
rawData = rawData.to(DEVICE)
contourData = contourData.to(DEVICE)
optimizer.zero_grad()
loss = F.l1_loss(model(rawData), contourData)
loss.backward()
optimizer.step()
print("[Iteration %d] [loss: %f]" % (iteration, loss.item()))
if iteration % 100 == 0:
torch.save(model.state_dict(),
os.path.join('weight', '{}.pth'.format(iteration)))
if iteration < 4000:
scheduler.step()
def __init__(self):
self.DEVICE = torch.device("cuda" if config.is_cuda else "cpu")
dataset = PairDataset(config.data_path,
max_src_len=config.max_src_len,
max_tgt_len=config.max_tgt_len,
truncate_src=config.truncate_src,
truncate_tgt=config.truncate_tgt)
self.vocab = dataset.build_vocab(embed_file=config.embed_file)
self.model = Seq2seq(self.vocab)
self.stop_word = list(
set([
self.vocab[x.strip()]
for x in open(config.stop_word_file).readlines()
]))
self.model.load_model()
self.model.to(self.DEVICE)
def __init__(self):
self.DEVICE = config.DEVICE
dataset = PairDataset(config.data_path,
max_src_len=config.max_src_len,
max_tgt_len=config.max_tgt_len,
truncate_src=config.truncate_src,
truncate_tgt=config.truncate_tgt)
self.vocab = dataset.build_vocab(embed_file=config.embed_file)
self.model = PGN(self.vocab)
self.stop_word = list(
set([
self.vocab[x.strip()] for x in open(
config.stop_word_file, encoding='utf-8').readlines()
]))
self.model.load_model()
self.model.to(self.DEVICE)
class PairDatasetTest(unittest.TestCase):
image_size = 224
batch_size = 16
dataset = PairDataset()
def test_dataset(self):
t0 = time()
batch = self.dataset[0]
self.assertEqual(type(batch[0]), np.ndarray)
self.assertEqual(type(batch[1]), np.int64)
print("Dataset: %.2fs" % (time() - t0))
def test_split(self):
t0 = time()
train_ids, test_ids = self.dataset.get_train_test_split()
self.assertEqual(
train_ids.size(0) + test_ids.size(0), len(self.dataset))
self.dataset.transform = None
train_loader = thd.DataLoader(
self.dataset,
batch_size=1,
sampler=thd.SubsetRandomSampler(train_ids))
test_loader = thd.DataLoader(self.dataset,
batch_size=1,
sampler=thd.SubsetRandomSampler(test_ids))
loader_train_ids, _ = torch.tensor(
[batch[2][0] for batch in train_loader]).sort()
loader_test_ids, _ = torch.tensor(
[batch[2][0] for batch in test_loader]).sort()
self.assertEqual(
loader_train_ids.eq(train_ids.sort()[0]).sum(), train_ids.size(0))
self.assertEqual(
loader_test_ids.eq(test_ids.sort()[0]).sum(), test_ids.size(0))
# Check that test IDs do not leak with training
self.assertEqual(
np.intersect1d(loader_train_ids.numpy(),
test_ids.numpy()).shape[0], 0)
self.assertEqual(
np.intersect1d(loader_train_ids.numpy(),
loader_test_ids.numpy()).shape[0], 0)
self.assertEqual(
np.intersect1d(loader_test_ids.numpy(),
train_ids.numpy()).shape[0], 0)
print("Split: %.2fs" % (time() - t0))
default='./datasets/DUTS_TEST')
parser.add_argument('--cuda',
help="'cuda' for cuda, 'cpu' for cpu, default = cuda",
default='cuda',
choices=['cuda', 'cpu'])
parser.add_argument('--batch_size',
help="batchsize, default = 8",
default=8,
type=int)
args = parser.parse_args()
# TODO : Add multiGPU Model
device = torch.device(args.cuda)
batch_size = args.batch_size
test_dataset = PairDataset(args.dataset,
train=False,
data_augmentation=False)
load = args.load
model = Unet(cfg).to(device)
assert load is not None
state_dict = torch.load(load, map_location=device)
start_iter = int(load.split('epo_')[1].strip('step.ckpt'))
start_epo = int(load.split('/')[3].split('epo')[0])
now = datetime.datetime.strptime(load.split('/')[2], '%m%d%H%M')
print("Loading Model from {}".format(load))
print("Start_iter : {}".format(start_iter))
print("now : {}".format(now.strftime('%m%d%H%M')))
model.load_state_dict(state_dict)
source_transform = transform.Compose([
# RandomGaussianNoise(p=0.95, mean=0, std=25, fixed_distribution=False),
RandomTextOverlay(p=1, max_occupancy=30, length=(15, 30)),
transform.ToTensor(),
])
target_transform = transform.Compose([
# RandomGaussianNoise(p=0.95, mean=0, std=25, fixed_distribution=False),
RandomTextOverlay(p=1, max_occupancy=30, length=(15, 30)),
transform.ToTensor(),
])
test_transform = transform.ToTensor()
train_set = PairDataset(root_dir=os.path.join(args.data_path, 'train'), pre_transform=pre_transform,
source_transform=source_transform, target_transform=target_transform)
test_set = PairDataset(root_dir=os.path.join(args.data_path, 'test'), pre_transform=pre_transform,
source_transform=source_transform, target_transform=test_transform)
train_data = DataLoader(train_set, batch_size, shuffle=True, pin_memory=True, num_workers=num_workers, drop_last=True)
test_data = DataLoader(test_set, batch_size, shuffle=False, pin_memory=True, num_workers=num_workers, drop_last=False)
model = UNet()
# model = ResNoise()
KaimingInitializer(model)
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=lr, betas=adam_param[0:2], eps=adam_param[2],
weight_decay=args.wd)
lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=epochs / 5, factor=0.5, verbose=True)
model = nn.DataParallel(model)
parser.add_argument('--epoch', help='# of epochs. default = 20', default=20, type=int)
parser.add_argument('-lr', '--learning_rate', help='learning_rate. default = 0.001', default=0.001, type=float)
parser.add_argument('--lr_decay', help='Learning rate decrease by lr_decay time per decay_step, default = 0.1',
default=0.1, type=float)
parser.add_argument('--decay_step', help='Learning rate decrease by lr_decay time per decay_step, default = 7000',
default=7000, type=int)
parser.add_argument('--display_freq', help='display_freq to display result image on Tensorboard',
default=1000, type=int)
args = parser.parse_args()
# TODO : Add multiGPU Model
device = torch.device(args.cuda)
batch_size = args.batch_size
epoch = args.epoch
duts_dataset = PairDataset(args.dataset)
load = args.load
start_iter = 0
model = Unet(cfg).cuda()
#vgg = torchvision.models.vgg16(pretrained=True)
#model.encoder.seq.load_state_dict(vgg.features.state_dict())
now = datetime.datetime.now()
start_epo = 0
#del vgg
if load is not None:
state_dict = torch.load(load, map_location=args.cuda)
#start_iter = int(load.split('epo_')[1].strip('step.ckpt')) + 1
#start_epo = int(load.split('/')[3].split('epo')[0])
#now = datetime.datetime.strptime(load.split('/')[2], '%m%d%H%M')
############
if train_backbone:
optimizer = torch.optim.Adam(net.parameters(),
betas=(0.9, 0.98),
eps=1e-06,
weight_decay=0.,
lr=1e-05)
else:
optimizer = torch.optim.RMSprop(list(net.vis_transform.parameters()) +
list(net.txt_transform.parameters()),
lr=0.0001)
# Datasets
###########
ds = PairDataset(
"./data/tgif-msrvtt10k/FeatureData/mean_resnext101_resnet152/",
"./data/tgif-msrvtt10k/TextData/tgif-msrvtt10k.caption.txt")
ds = RobertaDataset(ds,
net.roberta.encode,
bow_vocab_file=vocab_file,
w2v_file=w2v_file)
data_loader = torch.utils.data.DataLoader(dataset=ds,
batch_size=batch_size,
shuffle=True,
pin_memory=False,
num_workers=4,
collate_fn=roberta_collate_tokens)
val_ds = PairDataset(
"./data/tv2016train/FeatureData/mean_resnext101_resnet152/",
"./data/tv2016train/TextData/setA/tv2016train.caption.txt")
def generate_self_train_translations(
train_dataset: PairDataset,
model: Model,
chemprop_predictor: ChempropPredictor,
args: Namespace,
k: int = 4) -> Tuple[List[List[str]], List[List[float]]]:
# Copy train dataset since we're going to modify it
train_dataset = deepcopy(train_dataset)
# Keep a copy of the original targets and their properties for molecules which don't get enough translations
src = deepcopy(train_dataset.src)
tgt = deepcopy(train_dataset.tgt)
if not (hasattr(train_dataset, 'tgt_props')
and train_dataset.tgt_props is not None):
train_dataset.tgt_props = chemprop_predictor(train_dataset.tgt_smiles)
tgt_props = deepcopy(train_dataset.tgt_props)
# Increase batch size since using no_grad
train_dataset.set_batch_size(train_dataset.batch_size * 4)
# Initialize mapping from index in train_dataset.src to index in all_translations/props
# as train_dataset.src gets filtered
index_mapping = list(range(len(train_dataset)))
if args.unconditional:
for batch_src, _ in train_dataset: # get an arbitrarily batch_src as a dummy
# Pad batch src
batch_src, lengths_src = pad_batch(batch_src,
train_dataset.pad_index)
# To cuda
batch_src, lengths_src = batch_src.cuda(), lengths_src.cuda()
break
if args.dupe:
good_translations_and_props = []
else:
good_translations_and_props = set()
for _ in range(args.num_translations_per_input):
if args.dupe:
translations = []
else:
translations = set()
translations.update(train_dataset.tgt_smiles)
for _ in range(
math.ceil(
len(train_dataset) /
train_dataset.batch_size)): # keep same num attempts
# Predict
with torch.no_grad():
batch_indices = model.predict(batch_src,
lengths_src,
sample=True) # seqlen x bs
# Transpose
batch_indices = torch.stack(batch_indices,
dim=1).tolist() # bs x seqlen
# Convert predicted indices to smiles
batch_translations = [
train_dataset.indices2smiles(indices)
for indices in batch_indices
]
if args.dupe:
translations += batch_translations
else:
translations.update(batch_translations)
if not args.dupe:
translations = translations.difference(
train_dataset.tgt_smiles) # remove the ones we already had
translations = list(translations)
props = chemprop_predictor(translations)
chemprop_predictor.clear_cache()
for t, prop in zip(translations, props):
if prop is not None and prop > args.prop_min and (
args.prop_max is None or prop <= args.prop_max):
if args.dupe:
good_translations_and_props.append((t, prop))
else:
good_translations_and_props.add((t, prop))
if len(good_translations_and_props) >= k * len(train_dataset):
break
good_translations_and_props = list(good_translations_and_props)
random.shuffle(good_translations_and_props)
good_translations_and_props = good_translations_and_props[:k * len(
train_dataset)]
return [x[0] for x in good_translations_and_props
], [x[1] for x in good_translations_and_props
] # translations, props
# Initialize set of indices which have k translations
finished_indices = set()
# Generate translations
all_translations_and_props = [set() for _ in range(len(train_dataset))]
for _ in range(args.num_translations_per_input):
translations = []
translation_index_batches = []
for batch_src, _ in tqdm(train_dataset,
total=math.ceil(
len(train_dataset) /
train_dataset.batch_size)):
# Pad batch src
batch_src, lengths_src = pad_batch(batch_src,
train_dataset.pad_index)
# To cuda
batch_src, lengths_src = batch_src.cuda(), lengths_src.cuda()
# Predict
with torch.no_grad():
batch_indices = model.predict(batch_src,
lengths_src,
sample=True) # seqlen x bs
# Transpose
batch_indices = torch.stack(batch_indices,
dim=1).tolist() # bs x seqlen
# Convert predicted indices to smiles
batch_translations = [
train_dataset.indices2smiles(indices)
for indices in batch_indices
]
# Extend
translations += batch_translations
translation_index_batches.append(batch_indices)
# Predict properties
props = chemprop_predictor(translations)
chemprop_predictor.clear_cache()
assert len(train_dataset) == len(translations) == len(props)
# Find indices of good translations; if source side, we're checking tgt smiles with the new translations for similarity
good_translation_indices = find_good_translations(
src_smiles=train_dataset.src_smiles,
translations=translations,
args=args,
props=props,
tgt_props=tgt_props,
prop_min=args.prop_min,
prop_max=args.prop_max,
sim_threshold=args.morgan_similarity_threshold,
translation_index_batches=translation_index_batches)
# Add good translations to all translations
for i in good_translation_indices:
j = index_mapping[
i] # convert from index in current train_dataset.src to index in all_translations/props
# Add good translation
all_translations_and_props[j].add((translations[i], props[i]))
# Indicate if we have enough translations for a molecule
if len(all_translations_and_props[j]) >= k:
finished_indices.add(j)
# Update index_mapping to only keep src/tgt molecules without enough good translations
keep_indices = [
i for i in range(len(translations))
if index_mapping[i] not in finished_indices
]
train_dataset.src = [train_dataset.src[i] for i in keep_indices]
train_dataset.src_smiles = [
train_dataset.src_smiles[i] for i in keep_indices
]
train_dataset.tgt = [train_dataset.tgt[i] for i in keep_indices]
train_dataset.tgt_smiles = [
train_dataset.tgt_smiles[i] for i in keep_indices
]
train_dataset.tgt_props = [
train_dataset.tgt_props[i] for i in keep_indices
]
index_mapping = [index_mapping[i] for i in keep_indices]
# Check lengths
assert len(all_translations_and_props) == len(tgt) == len(tgt_props)
# Add actual target to any molecules which did not get enough translations
all_translations, all_props = [], []
for i in range(len(all_translations_and_props)):
all_translations_and_props[i] = list(all_translations_and_props[i])
all_translations_and_props[i] = [
(train_dataset.smiles2indices(translation), prop)
for translation, prop in all_translations_and_props[i]
]
if len(all_translations_and_props[i]) < k:
all_translations_and_props[i] += [
(tgt[i], tgt_props[i])
] * (k - len(all_translations_and_props[i]))
all_translations.append(
[pair[0] for pair in all_translations_and_props[i]])
all_props.append([pair[1] for pair in all_translations_and_props[i]])
return all_translations, all_props
# Update minimum evaluating loss.
if (avg_val_loss < val_losses):
torch.save(model.encoder, config.encoder_save_name)
torch.save(model.decoder, config.decoder_save_name)
torch.save(model.attention, config.attention_save_name)
torch.save(model.reduce_state, config.reduce_state_save_name)
val_losses = avg_val_loss
with open(config.losses_path, 'wb') as f:
pickle.dump(val_losses, f)
writer.close()
if __name__ == "__main__":
# Prepare dataset for training.
DEVICE = torch.device('cuda') if config.is_cuda else torch.device('cpu')
dataset = PairDataset(config.data_path,
max_src_len=config.max_src_len,
max_tgt_len=config.max_tgt_len,
truncate_src=config.truncate_src,
truncate_tgt=config.truncate_tgt)
val_dataset = PairDataset(config.val_data_path,
max_src_len=config.max_src_len,
max_tgt_len=config.max_tgt_len,
truncate_src=config.truncate_src,
truncate_tgt=config.truncate_tgt)
vocab = dataset.build_vocab(embed_file=config.embed_file)
train(dataset, val_dataset, vocab, start_epoch=0)
# -*- coding: utf-8 -*-
from transform import RandomTextOverlay, RandomCrop, RandomGaussianNoise
from dataset import PairDataset
import cv2
import os
dr_root = '/media/piston/data/Noise2Noise/'
pr = RandomCrop(256, pad_if_needed=True)
# ts = RandomTextOverlay(p=1, max_occupancy=50, length=(15, 30))
ts = RandomGaussianNoise(p=0.95, mean=0, std=25, fixed_distribution=False)
train_set = PairDataset(root_dir=os.path.join(dr_root, 'train'),
pre_transform=pr,
source_transform=ts,
target_transform=ts)
for s, t in train_set:
print(s.shape, t.shape)
cv2.imshow('s', s)
cv2.imshow('t', t)
cv2.waitKey(0)
def split(root_dir, target_dir):
import os
import glob
import random
import shutil
imgs = glob.glob(os.path.join(root_dir, '*.JPEG'))
type=int)
parser.add_argument('--cont',
help="Measure scores from this iter",
default=0,
type=int)
parser.add_argument('--step',
help="Measure scores per this iter step",
default=10000,
type=int)
args = parser.parse_args()
models = sorted(os.listdir(args.model_dir),
key=lambda x: int(x.split('epo_')[1].split('step')[0]))
pairdataset = PairDataset(root_dir=args.dataset,
train=False,
data_augmentation=False)
dataloader = DataLoader(pairdataset, 8, shuffle=True)
beta_square = 0.3
device = torch.device("cuda")
if args.logdir is not None:
writer = SummaryWriter(args.logdir)
model = Unet().to(device)
for model_name in models:
print('model', model_name)
model_iter = int(model_name.split('epo_')[1].split('step')[0])
#if model_iter % args.step != 0:
# continue
#if model_iter < args.cont:
# continue
#if args.which_iter > 0 and args.which_iter != model_iter:
args = get_args()
dataset = load_pkl("dataset.pkl")
raw_set = dataset['raw_set']
train_set = dataset['train_set']
test_set = dataset['test_set']
train_pair = dataset['train_pair']
user_size = dataset['user_size']
item_size = dataset['item_size']
if args.train:
dim = args.dim
batch_size = args.batch_size
dataset = PairDataset(item_size, train_set, train_pair)
dataloader = DataLoader(dataset, batch_size, num_workers=4)
model = Model(user_size, item_size, dim).cuda()
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-1)
max_score = 0
for epoch in range(args.epochs):
for idx, (u, i, j) in enumerate(dataloader):
opt.zero_grad()
x_ui, x_uj = model(u, i, j)
loss = bce_loss(x_ui, x_uj) if args.bce else bpr_loss(x_ui, x_uj)
loss.backward()
opt.step()
print(f"{idx}:{loss/batch_size:.4f}", end='\r')
result = get_top_k(model.W.detach(),
model.H.detach(),
def main(args: Namespace):
if args.unconditional:
assert args.morgan_similarity_threshold == 0 # shouldn't care about inputs in this case
i2s = None
if args.checkpoint_dir is not None:
assert args.checkpoint_path is None
for _, _, files in os.walk(args.checkpoint_dir):
for fname in files:
if fname.endswith('.pt'):
args.checkpoint_path = os.path.join(
args.checkpoint_dir, fname)
if args.checkpoint_path is not None:
print('loading model from checkpoint')
model, i2s = load_model(args)
full_train_dataset = PairDataset(
path=args.train_path,
i2s=i2s,
batch_size=args.batch_size,
extra_vocab_path=args.extra_precursors_path
if args.extra_precursors_path is not None else None,
max_data=args.train_max_data
if args.train_max_data is not None else None)
pair_datasets = full_train_dataset.split([0.9, 0.1], seed=0)
train_dataset, val_dataset = pair_datasets[0], pair_datasets[1]
predict_dataset = SourceDataset(path=args.val_path,
i2s=train_dataset.i2s,
s2i=train_dataset.s2i,
pad_index=train_dataset.pad_index,
start_index=train_dataset.start_index,
end_index=train_dataset.end_index,
batch_size=args.batch_size)
if args.checkpoint_path is None:
print('building model from scratch')
model = Model(args=args,
vocab_size=len(train_dataset.i2s),
pad_index=train_dataset.pad_index,
start_index=train_dataset.start_index,
end_index=train_dataset.end_index)
for param in model.parameters():
if param.dim() == 1:
nn.init.constant_(param, 0)
else:
nn.init.xavier_normal_(param)
print(model)
print('num params: {:,}'.format(
sum(p.numel() for p in model.parameters() if p.requires_grad)))
model = model.cuda()
chemprop_predictor = ChempropPredictor(args)
criterion = LossFunction(train_dataset.pad_index, args.kl_weight)
optimizer = optim.Adam(model.parameters(), lr=args.init_lr)
scheduler = lr_scheduler.ExponentialLR(optimizer, 0.9)
for epoch in range(args.epochs):
print('epoch {}'.format(epoch))
train_dataset.reshuffle(seed=epoch)
train(model=model,
train_dataset=train_dataset,
criterion=criterion,
optimizer=optimizer,
max_grad_norm=args.max_grad_norm)
val_loss = validate(model=model,
val_dataset=val_dataset,
criterion=criterion)
os.makedirs(os.path.join(args.save_dir, 'epoch' + str(epoch)),
exist_ok=True)
train_dataset.save(
os.path.join(args.save_dir, 'epoch' + str(epoch),
'train_pairs.csv'))
save_model(model=model,
i2s=train_dataset.i2s,
path=os.path.join(args.save_dir, 'epoch' + str(epoch),
'val_loss_{}.pt'.format(val_loss)))
predict(model=model,
predict_dataset=predict_dataset,
save_dir=os.path.join(args.save_dir, 'epoch' + str(epoch)),
args=args,
chemprop_predictor=chemprop_predictor
if not args.no_predictor_at_val else None,
sample=not args.greedy_prediction,
num_predictions=args.val_num_predictions,
print_filter_frac=args.print_filter_frac)
if epoch % args.evaluate_every == 0:
evaluate(pred_smiles_dir=os.path.join(args.save_dir,
'epoch' + str(epoch)),
train_path=args.train_path,
val_path=args.val_path,
checkpoint_dir=args.chemprop_dir,
computed_prop=args.computed_prop,
prop_min=args.prop_min,
sim_thresholds=[0.2, 0.4, 0.6, 0.8, 0.9, 1.0],
chemprop_predictor=chemprop_predictor,
prop_max=args.prop_max,
unconditional=args.unconditional)
scheduler.step()
if args.self_train_epochs > 0:
# store parameters of current model for a loss to constrain it not to stray too far
original_parameter_vector = parameters_to_vector(
model.parameters()).data
parameter_crit = nn.MSELoss()
args.epoch_length = len(train_dataset.src) // 2
# Get properties of target molecules in train set
train_dataset.tgt_props = np.array(
chemprop_predictor(train_dataset.tgt_smiles))
augmented_train_dataset = deepcopy(train_dataset)
epochs_to_dataset_creation = 0
for epoch in range(args.epochs, args.epochs + args.self_train_epochs):
print('self train epoch {}'.format(epoch))
if epochs_to_dataset_creation == 0:
train_dataset.reshuffle(seed=epoch)
if args.self_train_max_data is not None:
self_train_dataset = deepcopy(train_dataset)
self_train_dataset.src, self_train_dataset.tgt = \
self_train_dataset.src[:args.self_train_max_data], self_train_dataset.tgt[:args.self_train_max_data]
self_train_dataset.src_smiles, self_train_dataset.tgt_smiles = \
self_train_dataset.src_smiles[:args.self_train_max_data], self_train_dataset.tgt_smiles[:args.self_train_max_data]
if hasattr(self_train_dataset, 'src_props'):
self_train_dataset.src_props = self_train_dataset.src_props[:
args
.
self_train_max_data]
if hasattr(self_train_dataset, 'tgt_props'):
self_train_dataset.tgt_props = self_train_dataset.tgt_props[:
args
.
self_train_max_data]
else:
self_train_dataset = deepcopy(train_dataset)
if args.extra_precursors_path is not None:
self_train_dataset.add_dummy_pairs(
args.extra_precursors_path)
translations, props = generate_self_train_translations(
train_dataset=self_train_dataset,
model=model,
chemprop_predictor=chemprop_predictor,
args=args,
k=args.k)
if not args.keep_translations: # drop old translations and restart
augmented_train_dataset = deepcopy(self_train_dataset)
if args.unconditional:
new_train_dataset = deepcopy(self_train_dataset)
new_train_dataset.tgt_smiles = translations
new_train_dataset.tgt = [
list(self_train_dataset.smiles2indices(smiles))
for smiles in new_train_dataset.tgt_smiles
]
new_train_dataset.tgt = np.array(new_train_dataset.tgt)
new_train_dataset.src_smiles = translations # any dummy is fine
new_train_dataset.src = [
list(self_train_dataset.smiles2indices(smiles))
for smiles in new_train_dataset.src_smiles
]
new_train_dataset.src = np.array(new_train_dataset.src)
else:
new_train_dataset = deepcopy(self_train_dataset)
new_train_dataset.src = np.concatenate(
[self_train_dataset.src for _ in range(args.k)])
new_train_dataset.src_smiles = []
for _ in range(args.k):
new_train_dataset.src_smiles += self_train_dataset.src_smiles
new_train_dataset.tgt = []
for i in range(args.k):
new_train_dataset.tgt += [
translations[j][i]
for j in range(len(translations))
]
new_train_dataset.tgt_smiles = [
self_train_dataset.indices2smiles(indices)
for indices in new_train_dataset.tgt
]
new_train_dataset.tgt = np.array(new_train_dataset.tgt)
if args.replace_old_dataset:
augmented_train_dataset = new_train_dataset
else:
augmented_train_dataset.add(new_train_dataset)
if not args.unconditional:
augmented_train_dataset.filter_dummy_pairs(
need_props=False) # filters src == tgt pairs
epochs_to_dataset_creation = args.epochs_per_dataset
augmented_train_dataset.reshuffle(seed=epoch, need_props=False)
epochs_to_dataset_creation -= 1
train(model=model,
train_dataset=augmented_train_dataset,
criterion=criterion,
optimizer=optimizer,
max_grad_norm=args.max_grad_norm,
original_parameter_vector=original_parameter_vector,
parameter_crit=parameter_crit,
parameter_crit_weight=args.l2_diff_weight)
val_loss = validate(model=model,
val_dataset=val_dataset,
criterion=criterion)
os.makedirs(os.path.join(args.save_dir, 'epoch' + str(epoch)),
exist_ok=True)
augmented_train_dataset.save(
os.path.join(args.save_dir, 'epoch' + str(epoch),
'train_pairs.csv'))
save_model(model=model,
i2s=train_dataset.i2s,
path=os.path.join(args.save_dir, 'epoch' + str(epoch),
'val_loss_{}.pt'.format(val_loss)))
predict(model=model,
predict_dataset=predict_dataset,
save_dir=os.path.join(args.save_dir, 'epoch' + str(epoch)),
args=args,
chemprop_predictor=chemprop_predictor
if not args.no_predictor_at_val else None,
sample=not args.greedy_prediction,
num_predictions=args.val_num_predictions,
print_filter_frac=args.print_filter_frac)
evaluate(pred_smiles_dir=os.path.join(args.save_dir,
'epoch' + str(epoch)),
train_path=args.train_path,
val_path=args.val_path,
checkpoint_dir=args.chemprop_dir,
computed_prop=args.computed_prop,
prop_min=args.prop_min,
sim_thresholds=[0.2, 0.4, 0.6, 0.8, 0.9, 1.0],
chemprop_predictor=chemprop_predictor,
prop_max=args.prop_max,
unconditional=args.unconditional)
scheduler.step()
# for convenient evaluation
os.makedirs(os.path.join(args.save_dir, 'final_eval'), exist_ok=True)
test_dataset = SourceDataset(path=args.test_path,
i2s=train_dataset.i2s,
s2i=train_dataset.s2i,
pad_index=train_dataset.pad_index,
start_index=train_dataset.start_index,
end_index=train_dataset.end_index,
batch_size=args.batch_size)
predict(model=model,
predict_dataset=test_dataset,
save_dir=os.path.join(args.save_dir, 'final_eval'),
args=args,
chemprop_predictor=chemprop_predictor
if not args.no_predictor_at_val else None,
sample=not args.greedy_prediction,
num_predictions=args.val_num_predictions,
print_filter_frac=args.print_filter_frac)
if args.final_eval_chemprop_dir is not None:
args.computed_prop = None
args.chemprop_dir = args.final_eval_chemprop_dir
chemprop_predictor = ChempropPredictor(args)
if args.final_eval_computed_prop is not None:
args.chemprop_dir = None
args.computed_prop = args.final_eval_computed_prop
chemprop_predictor = ChempropPredictor(args)
evaluate(pred_smiles_dir=os.path.join(args.save_dir, 'final_eval'),
train_path=args.train_path,
val_path=args.test_path,
checkpoint_dir=args.chemprop_dir,
computed_prop=args.computed_prop,
prop_min=args.prop_min,
sim_thresholds=[0.2, 0.4, 0.6, 0.8, 0.9, 1.0],
chemprop_predictor=chemprop_predictor,
prop_max=args.prop_max,
unconditional=args.unconditional)
help=
'Learning rate decrease by lr_decay time per decay_step, default = 7000',
default=7000,
type=int)
parser.add_argument(
'--display_freq',
help='display_freq to display result image on Tensorboard',
default=1000,
type=int)
args = parser.parse_args()
# TODO : Add multiGPU Model
device = torch.device(args.cuda)
batch_size = args.batch_size
epoch = args.epoch
train_dataset = PairDataset(args.dataset)
load = args.load
start_iter = 0
model = Unet(cfg).cuda()
vgg = torchvision.models.vgg16(pretrained=True)
model.encoder.seq.load_state_dict(vgg.features.state_dict())
now = datetime.datetime.now()
start_epo = 0
del vgg
if load is not None:
state_dict = torch.load(load, map_location=args.cuda)
start_iter = int(load.split('epo_')[1].strip('step.ckpt'))
start_epo = int(load.split('/')[3].split('epo')[0])
now = datetime.datetime.strptime(load.split('/')[2], '%m%d%H%M')
# model = nn.DataParallel(model)
model = wapper(model)
model.load_state_dict(
torch.load('../param/5_gaussian.pt', map_location=torch.device('cpu')))
model = model.module.cpu()
pre_transform = RandomCrop(256, pad_if_needed=True)
source_transform = transform.Compose([
RandomGaussianNoise(p=0.95, mean=0, std=25, fixed_distribution=False),
# RandomTextOverlay(p=1, max_occupancy=30, length=(15, 30)),
transform.ToTensor(),
])
test_transform = transform.ToTensor()
dt = PairDataset('/media/piston/data/Noise2Noise/test',
pre_transform=pre_transform,
source_transform=source_transform,
target_transform=test_transform)
def get_psnr(input, target):
"""Computes peak signal-to-noise ratio."""
return 10 * torch.log10(1 / F.mse_loss(input, target))
def to_numpy(tensor):
tensor = tensor.squeeze().numpy()
tensor = tensor.transpose(1, 2, 0)
tensor = np.clip(tensor * 255, 0, 255).astype(np.uint8)[..., ::-1]
return tensor
