def main():
gpus = digit_counter(os.environ["CUDA_VISIBLE_DEVICES"])[0]
params = [[0.001, 80, 10, 2, 0.6], [0.001, 320, 40, 8, 0.6],
[0.001, 160, 20, 4, 0.6], [0.001, 640, 80, 16, 0.6],
[0.001, 960, 120, 24, 0.6], [0.001, 1280, 160, 32, 0.6]]
# [0.001, 1600, 200, 40, 0.6]]
for id, param in enumerate(params):
name = str(param) + str(datetime.now()).replace(" ", "_")
# tensorboard_cb = keras.callbacks.TensorBoard(log_dir='./graph', histogram_freq=0,
# write_graph=True, write_images=True)
stop = EarlyStopping(monitor='val_loss',
min_delta=0.001,
patience=5,
verbose=1,
mode='auto')
# checkpointer = ModelCheckpoint(filepath='../weights/checkpoint.hdf5', verbose=1,
# save_best_only=True)
history = LossHistory()
callbacks = [history, stop]
if USE_MULTI is True:
parallel_model, model = network.create_model(model_params=param,
multi_gpu=USE_MULTI,
gpus=gpus)
train_with_all(DATASET_PATH,
VALSET_PATH,
target_model_name=name,
model=parallel_model,
save_model=model,
nb_epochs=20,
callbacks=callbacks)
else:
model = network.create_model(model_params=param,
multi_gpu=USE_MULTI,
gpus=gpus)
train_with_all(DATASET_PATH,
VALSET_PATH,
target_model_name=name,
model=model,
save_model=model,
nb_epochs=20,
callbacks=callbacks)
printc("Param set: {} done".format(id))
printc("Params: {}".format(param))
printc("All done", 'okgreen')
def main():
batch_size = 16
# Prepare for creating the Batch Generator
list_id, folder_start_pos = prepare_for_loader()
# Create the Batch Generator for training Data
trainObj = DataGenerator(MAIN_DIR,
list_id,
folder_start_pos,
batch_size,
shuffle=True)
model = create_model(batch_size)
adam = keras.optimizers.Adam(lr=1e-3)
model.compile(loss='mean_absolute_error',
optimizer=adam,
metrics=['accuracy'])
# Fit the model
model.fit_generator(generator=trainObj,
epochs=50,
verbose=2,
use_multiprocessing=True)
# Save the weights for the model
model.save_weights('cnn_rnn_weights.hk5')
def run_cv(fold_iterator, logger, params_dict, upsample=True):
for traindirs, testdirs in fold_iterator:
# TRAIN LOCAL PREDICTION MODEL
# Generators
logger.info('############ FOLD #############')
logger.info('Training folders are {}'.format(traindirs))
training_generator = DataLoader(data_dir,
traindirs,
32,
width_template=params_dict['width'],
upsample=upsample)
validation_generator = DataLoader(data_dir,
testdirs,
32,
width_template=params_dict['width'],
type='val',
upsample=upsample)
# Design model
model = create_model(params_dict['width'] + 1,
params_dict['h1'],
params_dict['h2'],
params_dict['h3'],
embed_size=params_dict['embed_size'],
drop_out_rate=params_dict['dropout_rate'],
use_batch_norm=params_dict['use_batchnorm'])
# Train model on training dataset
'''
model.fit_generator(generator=training_generator,
validation_data=validation_generator,
use_multiprocessing=True,
epochs=params_dict['n_epochs'],
workers=6)
'''
try:
model.load_weights(os.path.join(checkpoint_dir, 'model22.h5'))
except OSError:
print('here')
model.fit_generator(generator=training_generator,
validation_data=validation_generator,
use_multiprocessing=True,
epochs=params_dict['n_epochs'],
workers=4,
max_queue_size=20)
model.save_weights(os.path.join(checkpoint_dir, 'model.h5'))
metrics = model.evaluate_generator(generator=validation_generator,
workers=4,
max_queue_size=20)
logger.info(metrics)
def main():
test_model = create_model()
test_model.load_weights('cnn_rnn_weights.hk5')
files = sorted(glob.glob(TESTING_DATA_FRAMES_PATH))
for i in range(1, len(files)):
X = np.empty((1, 1, 240, 320, 3))
image = cv2.imread(files[i])
for j in range(1):
X[0, j, ] = cv2.imread(files[i - 1 + j])
bbox = test_model.predict(X)
bbox = np.int32(bbox[0])
cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]),
(0, 255, 0), 2)
cv2.imwrite(MAIN_DIR + "/predictions/{0}.jpg".format(i), image)
depthdata=util.read_cached_data(64)['depth']
train_iter = util.create_iterator(
data,
BATCH_SIZE)
test_iter=util.create_iterator(test_data,BATCH_SIZE,istrain=False)
train_loader = torch.utils.data.DataLoader(
train_iter,
batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=4)
test_loader = torch.utils.data.DataLoader(
test_iter,
batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=4)
#==========================================================================
net = network.create_model()
optimizer = optim.SGD(net.parameters(), lr=0.00001, momentum=0.9)
scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)
best_path = os.path.join(CHECK_DIR, 'best.ckpt')
if os.path.isfile(best_path):
checkpoint = torch.load(best_path,map_location=device)
global_t=checkpoint['global_t']
net.load_state_dict(checkpoint['state_dict'],strict=False)
net.train()
print("=> loaded checkpoint '{}' (global_t {})"
.format(best_path, checkpoint['global_t']))
else:
global_t=0
net.train()
print("=> no checkpoint found at '{}'".format(best_path))
import os
import random
import numpy as np
from options import parse_option
from network import create_model
from utils.core import evaluate
from datasets import get_dataloader
import warnings
warnings.filterwarnings('ignore')
# Seed
torch.manual_seed(0)
torch.cuda.manual_seed(0)
np.random.seed(0)
random.seed(0)
#NOTE: main loop for training
if __name__ == "__main__":
# Option
opt = parse_option(print_option=False)
# Data Loader
_, dataset_val = get_dataloader(opt)
# Network
net = create_model(opt)
# Evaluate
evaluate(dataset_val, net, opt)
import numpy as np
from network import create_model
from threading import Thread
from json import loads, dumps
from kafka import KafkaConsumer, KafkaProducer
from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
filename = 'weights-ep2076-val_loss0.2518.hdf5'
net = create_model(training=False)
net.load_weights(f'/tmp/src/Ascalon.NeuralNetwork.Service/ckpt/{filename}')
producer = KafkaProducer(bootstrap_servers=['35.189.215.83:9092'],
value_serializer=lambda x: dumps(x).encode('utf-8'))
consumer = KafkaConsumer(
'neuralnetwork_data',
bootstrap_servers=['35.189.215.83:9092'],
auto_offset_reset='earliest',
enable_auto_commit=True,
partition_assignment_strategy=[RoundRobinPartitionAssignor],
group_id='NeuralNetworkService',
value_deserializer=lambda x: loads(x.decode('utf-8')))
from pathlib import Path
import shutil
import torch.nn.functional as F
import testenv as scene_loader
import multiprocessing as mp
import threading as td
import time
import sys
mode='eva'
CHECK_DIR='./checkpoint'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
with open ('updateevafeaturedes', 'rb') as ft:
wholeData=pickle.load(ft)
#=====================================================================================================
net = network.create_model(pretrained=False)
best_path = os.path.join(CHECK_DIR, 'best.ckpt') #
if os.path.isfile(best_path):
checkpoint = torch.load(best_path)
global_t=checkpoint['global_t']
net.load_state_dict(checkpoint['state_dict'],strict=False)
print("=> loaded checkpoint '{}' (global_t {})"
.format(best_path, checkpoint['global_t']))
else:
print("can't evaluate!!!!!")
worlds= ['Home_011_1']
ACTIONS=['right', 'rotate_cw', 'rotate_ccw', 'forward', 'left', 'backward', 'stop']
#======================================================================================================
env =scene_loader.ActiveVisionDatasetEnv(world=worlds[0])
def getImagedes(world,startid):
importlib.reload(network)
importlib.reload(loss)
importlib.reload(training)
type1 = 'Water'
type2 = 'Grass'
save_path = './pokemon_results_' + type1 + '_' + type2 + '/'
# Get data
Xdset = dataloader.PokemonData(type=type1, data_path='./images_with_types')
dataloader_X = DataLoader(Xdset, batch_size=8, shuffle=True)
Ydset = dataloader.PokemonData(type=type2, data_path='./images_with_types')
dataloader_Y = DataLoader(Ydset, batch_size=8, shuffle=True)
# Get the model and train
G_XtoY, G_YtoX, D_X, D_Y = network.create_model()
losses, G_XtoY, G_YtoX, D_X, D_Y = training.training_loop(
dataloader_X=dataloader_X,
dataloader_Y=dataloader_Y,
n_epochs=4000,
G_XtoY=G_XtoY,
G_YtoX=G_YtoX,
D_X=D_X,
D_Y=D_Y,
lr=0.0005,
beta1=0.5,
beta2=0.999,
save_path=save_path)
# save the model
checkpoint_dir = './Checkpoints/'
file_source_lang='small_vocab_en',
file_target_lang='small_vocab_fr',
load_method=load_from_file,
)
# create mappings
x_sent_array, x_vocab_len, x_word_to_idx, x_idx_to_word = preprocess_data(
x, conf['MAX_LEN'], conf['VOCAB_SIZE'], conf['NUM_SENT'])
y_sent_array, y_vocab_len, y_word_to_idx, y_idx_to_word = preprocess_data(
y, conf['MAX_LEN'], conf['VOCAB_SIZE'], conf['NUM_SENT'])
# Find the length of the longest sequence
x_max_len = max([len(sentence) for sentence in x_sent_array])
y_max_len = max([len(sentence) for sentence in y_sent_array])
# Padding zeros to make all sequences have a same length with the longest one
print('Zero padding...')
X = pad_sequences(x_sent_array, maxlen=x_max_len, dtype='int32')
y = pad_sequences(y_sent_array, maxlen=y_max_len, dtype='int32')
# Creating the network model
print('Compiling model...')
model = create_model(x_vocab_len, x_max_len, y_vocab_len, y_max_len,
conf['HIDDEN_DIM'], conf['LAYER_NUM'])
# Finding trained weights of previous epoch if any
saved_weights = find_checkpoint_file('.')
saved_weights = []
train(X, y, y_word_to_idx, y_max_len, saved_weights, model, conf)
x_train = []
y_train = []
image_class = -1
for person in people:
folder = "/".join([data_dir, person])
image_class += 1
for image in os.listdir(folder):
image = "/".join([folder, image])
image = cv2.imread(image)
image = cv2.resize(image, (WIDTH, HEIGHT))
x_train.append(image)
y_train.append(
[1 if i == image_class else 0 for i in range(class_number)])
x_train = np.array(x_train)
y_train = np.array(y_train)
model = create_model(WIDTH, HEIGHT, class_number)
model.fit(x_train,
y_train,
epochs=30,
batch_size=32,
validation_split=0.15,
shuffle=True)
model.save('model_state.pt')
import network
print("pobieranie danych z pliku:")
data = network.read_data()
print("podział danych na 4 wymagane zbiory:")
train_data, train_labels, test_data, test_labels = network.separate_sets(data)
print("przefiltrowanie tylko interesujących nas cech:")
selected_attributes = [2, 4]
train_data = network.filter_only_selected_attributes(train_data,
selected_attributes)
test_data = network.filter_only_selected_attributes(test_data,
selected_attributes)
print("stworzenie modelu z domyslnymi parametrami:")
model = network.create_model(attributes=len(selected_attributes))
print("uczenie:")
model.fit(x=train_data, y=train_labels, epochs=20, batch_size=32)
print("testowanie:")
evaluate = model.evaluate(x=test_data, y=test_labels)
print("wynik:")
print(evaluate[1])
# Przykład porównywania sieci dla różnej liczby epok uczenia
result = numpy.zeros(shape=[30, 2])
for x in range(30):
model = network.create_model(attributes=len(selected_attributes))
model.fit(x=train_data, y=train_labels, epochs=x, batch_size=32)
evaluate = model.evaluate(x=test_data, y=test_labels)
def predict(input_path, output_path, resources_path):
"""
This is the skeleton of the prediction function.
The predict function will build your model, load the weights from the checkpoint and write a new file (output_path)
with your predictions in the BIES format.
The resources folder should contain everything you need to make the predictions. It is the "resources" folder in your submission.
N.B. DO NOT HARD CODE PATHS IN HERE. Use resource_path instead, otherwise we will not be able to run the code.
:param input_path: the path of the input file to predict.
:param output_path: the path of the output file (where you save your predictions)
:param resources_path: the path of the resources folder containing your model and stuff you might need.
:return: None
"""
true_tags = os.path.join(resources_path, "dataset_new/predict/true_tags.txt")
input_path = os.path.join(cwd, resources_path, input_path)
output_path = os.path.join(cwd, resources_path, output_path)
true_tags_path = os.path.join(cwd, true_tags)
src = os.path.join(cwd, "%sdataset_new/all/datasetOutput/unique_unigrams_char_to_id.json"%resources_path)
dst = os.path.join(cwd, "%sdataset_new/predict/datasetOutput/unique_unigrams_char_to_id.json"%resources_path)
copyfile(src, dst)
VOCAB_SIZE = get_vsize('unigrams')
EMBEDDING_SIZE = 32
HIDDEN_SIZE = 100
idText = []
ckpt = os.path.join(cwd, "%sweights.hdf5"%resources_path)
model = create_model(VOCAB_SIZE, EMBEDDING_SIZE, HIDDEN_SIZE)
with open(input_path, 'rb') as fb:
raw_text = fb.read().decode('UTF-8')
lines = []
chars_file = []
# phase = 'training'
with open(input_path, 'rb') as f:
content = f.readline().decode('UTF-8')
while content:
lines.append(content)
content = f.readline().decode('UTF-8')
_, tags = generate(lines, 'predict')
idText = char_to_id('predict', lines, 'unigrams')
# print(idText)
idText = pad_sequences(idText, truncating='pre', padding='post', maxlen=MAX_LENGTH)
labels_pred = model.predict(np.array(idText))
label_i = np.argmax(labels_pred, axis=2)
prediction = []
labels_dict = {1:'B', 2:'I', 3:'E', 4:'S'}
# tags2 = []
with tqdm(desc="prediction file", total=len(idText)) as pbar:
for i in range(len(label_i)):
pbar.update(1)
sentence = label_i[i]
labels = []
num_char = np.count_nonzero(idText[i])
for char in sentence[0:num_char]:
if (char != 0):
labels.append(labels_dict[char])
prediction.append(labels)
# Write the predicted output
with open(output_path, "w") as p_t:
for p_tag in prediction:
for e in p_tag:
p_t.write(''.join(str(e)))
p_t.write('\n')
# write the correct tags file
with open(true_tags_path, "w") as t_t:
for t_tag in tags:
for e in t_tag:
t_t.write(''.join(str(e)))
t_t.write('\n')
pred_list = textToList(true_tags_path, 'labels')
true_list = textToList(output_path, 'labels')
for i in range(len(pred_list)):
if (len(pred_list[i]) != len(true_list[i])):
diff = len(pred_list[i])-len(true_list[i])
del pred_list[i][-diff:]
score(pred_list, true_list, verbose=True)
def predict(testdirs,
checkpoint_dir,
data_dir,
params_dict,
upsample=False,
resolution_df=None):
model = create_model(params_dict['width'] + 1,
params_dict['h1'],
params_dict['h2'],
params_dict['h3'],
embed_size=params_dict['embed_size'],
drop_out_rate=params_dict['dropout_rate'],
use_batch_norm=params_dict['use_batchnorm'])
model.load_weights(os.path.join(checkpoint_dir, 'model.h5'))
est_c1 = load(os.path.join(checkpoint_dir, 'est_c1.joblib'))
est_c2 = load(os.path.join(checkpoint_dir, 'est_c2.joblib'))
for testfolder in testdirs:
res_x, res_y = None, None
if upsample:
res_x, res_y = resolution_df.loc[resolution_df['scan'] ==
testfolder,
['res_x', 'res_y']].values[0]
annotation_dir = os.path.join(data_dir, testfolder, 'Annotation')
img_dir = os.path.join(data_dir, testfolder, 'Data')
list_imgs = [
os.path.join(img_dir, dI) for dI in os.listdir(img_dir)
if (dI.endswith('png') and not dI.startswith('.'))
]
list_label_files = [
dI for dI in os.listdir(annotation_dir)
if (dI.endswith('txt') and not dI.startswith('.'))
]
try:
img_init = np.asarray(
Image.open(os.path.join(img_dir, "{:04d}.png".format(1))))
except FileNotFoundError:
img_init = np.asarray(
Image.open(os.path.join(img_dir, "{:05d}.png".format(1))))
img_init = prepare_input_img(img_init, res_x, res_y, upsample)
for j, label_file in enumerate(list_label_files):
print(label_file)
img_current = img_init
df = pd.read_csv(os.path.join(annotation_dir, label_file),
header=None,
names=['id', 'x', 'y'],
sep='\s+')
if upsample:
df['x_newres'] = df['x'] * res_x / 0.4
df['y_newres'] = df['y'] * res_y / 0.4
else:
df['x_newres'] = df['x']
df['y_newres'] = df['y']
c1_init, c2_init = df.loc[df['id'] == 1,
['x_newres', 'y_newres']].values[0, :]
a, b = np.nonzero(img_init[:, 20:(len(img_init) - 20)])
if upsample:
list_centers = [[c1_init * 0.4 / res_x, c2_init * 0.4 / res_y]]
else:
list_centers = [[c1_init, c2_init]]
xax, yax = find_template_pixel(c1_init, c2_init,
params_dict['width'],
img_init.shape[1],
img_init.shape[0])
template_init = img_init[np.ravel(yax),
np.ravel(xax)].reshape(
1, len(yax), len(xax))
c1, c2 = c1_init, c2_init
k = 0
stop_temporal = False
for i in range(2, len(list_imgs) + 1):
if i % 100 == 0:
print(i)
try:
img_current = np.asarray(
Image.open(
os.path.join(img_dir, "{:04d}.png".format(i))))
except FileNotFoundError:
img_current = np.asarray(
Image.open(
os.path.join(img_dir, "{:05d}.png".format(i))))
img_current = prepare_input_img(img_current, res_x, res_y,
upsample)
if i > 5:
tmp = list_centers[-10:].reshape(-1, 2)
assert tmp.shape[0] == 5
c1, c2, stop_temporal, k = get_next_center(
k, stop_temporal, c1, c2, img_current, params_dict,
model, template_init, c1_init, c2_init, None, est_c1,
est_c2, tmp[:, 0], tmp[:, 1])
else:
c1, c2, stop_temporal, k = get_next_center(
k, stop_temporal, c1, c2, img_current, params_dict,
model, template_init, c1_init, c2_init, None)
# project back in init coords
if upsample:
c1_orig_coords = c1 * 0.4 / res_x
c2_orig_coords = c2 * 0.4 / res_y
else:
c1_orig_coords = c1
c2_orig_coords = c2
list_centers = np.append(list_centers,
[c1_orig_coords, c2_orig_coords])
list_centers = list_centers.reshape(-1, 2)
pred_df = pd.DataFrame()
pred_df['idx'] = range(1, len(list_centers) + 1)
pred_df['c1'] = list_centers[:, 0]
pred_df['c2'] = list_centers[:, 1]
pred_df.to_csv(os.path.join(checkpoint_dir,
'{}'.format(label_file)),
header=False,
index=False)
def train(traindirs,
data_dir,
upsample,
params_dict,
checkpointdir,
logger,
validation_gen=None):
if logger is not None:
logger.info('Training folders are {}'.format(traindirs))
else:
print('Training folders are {}'.format(traindirs))
training_generator = DataLoader(data_dir,
traindirs,
32,
width_template=params_dict['width'],
upsample=upsample)
earl = keras.callbacks.EarlyStopping(monitor='val_loss', patience=3)
# Design model
model = create_model(params_dict['width'] + 1,
params_dict['h1'],
params_dict['h2'],
params_dict['h3'],
embed_size=params_dict['embed_size'],
drop_out_rate=params_dict['dropout_rate'],
use_batch_norm=params_dict['use_batchnorm'])
# Train local Net
if validation_gen is None:
model.fit_generator(generator=training_generator,
use_multiprocessing=True,
epochs=params_dict['n_epochs'],
workers=4,
max_queue_size=20,
callbacks=[earl])
else:
model.fit_generator(generator=training_generator,
validation_data=validation_gen,
use_multiprocessing=True,
epochs=params_dict['n_epochs'],
workers=4,
max_queue_size=20)
if logger is not None:
logger.info('Local Net trained')
logger.info('Stopped epoch {}'.format(earl.stopped_epoch))
else:
print('Local Net trained')
print('Stopped epoch {}'.format(earl.stopped_epoch))
# Train the temporal model
for folder in traindirs:
if logger is not None:
logger.info('Getting temporal training set for {}'.format(folder))
else:
print('Getting temporal training set for {}'.format(folder))
img_dir = os.path.join(data_dir, folder, 'Data')
annotation_dir = os.path.join(data_dir, folder, 'Annotation')
list_label_files = [
os.path.join(annotation_dir, dI)
for dI in os.listdir(annotation_dir)
if (dI.endswith('txt') and not dI.startswith('.'))
]
try:
img_init = np.asarray(
Image.open(os.path.join(img_dir, "{:04d}.png".format(1))))
except FileNotFoundError:
img_init = np.asarray(
Image.open(os.path.join(img_dir, "{:05d}.png".format(1))))
list_imgs = [
os.path.join(img_dir, dI) for dI in os.listdir(img_dir)
if (dI.endswith('png') and not dI.startswith('.'))
]
n_obs = len(list_imgs)
X0, X1, X2, X3, X4, X5 = [], [], [], [], [], []
Y0, Y1, Y2, Y3, Y4, Y5 = [], [], [], [], [], []
for label in list_label_files:
print(label)
df = pd.read_csv(os.path.join(annotation_dir, label),
header=None,
names=['id', 'x', 'y'],
sep='\s+')
c1_interpolate = np.interp(np.arange(1, n_obs + 1), df.id.values,
df.x.values)
c2_interpolate = np.interp(np.arange(1, n_obs + 1), df.id.values,
df.y.values)
n = len(c1_interpolate)
X0 = np.append(X0, c1_interpolate)
X1 = np.append(X1, c1_interpolate[1:n])
X2 = np.append(X2, c1_interpolate[2:n])
X3 = np.append(X3, c1_interpolate[3:n])
X4 = np.append(X4, c1_interpolate[4:n])
X5 = np.append(X5, c1_interpolate[5:n])
Y0 = np.append(Y0, c2_interpolate)
Y1 = np.append(Y1, c2_interpolate[1:n])
Y2 = np.append(Y2, c2_interpolate[2:n])
Y3 = np.append(Y3, c2_interpolate[3:n])
Y4 = np.append(Y4, c2_interpolate[4:n])
Y5 = np.append(Y5, c2_interpolate[5:n])
l = len(X5)
fullX = np.transpose(
np.vstack([X0[0:l], X1[0:l], X2[0:l], X3[0:l], X4[0:l]]))
fullY = np.transpose(
np.vstack([Y0[0:l], Y1[0:l], Y2[0:l], Y3[0:l], Y4[0:l]]))
c1_label = X5
c2_label = Y5
est_c1 = RidgeCV()
est_c2 = RidgeCV()
scores_c1 = cross_validate(est_c1,
fullX,
c1_label,
cv=5,
scoring=('r2', 'neg_mean_squared_error'))
scores_c2 = cross_validate(est_c2,
fullY,
c2_label,
cv=5,
scoring=('r2', 'neg_mean_squared_error'))
if logger is not None:
logger.info('c1')
logger.info(scores_c1['test_neg_mean_squared_error'])
logger.info('c2')
logger.info(scores_c2['test_neg_mean_squared_error'])
else:
print('c1')
print(scores_c1['test_neg_mean_squared_error'])
print('c2')
print(scores_c2['test_neg_mean_squared_error'])
# Fit on the whole training set
est_c1.fit(fullX, c1_label)
est_c2.fit(fullY, c2_label)
# Save the local Net and the temporal model
if logger is not None:
logger.info('Saving trained models to {}'.format(checkpoint_dir))
else:
print('Saving trained models to {}'.format(checkpoint_dir))
model.save_weights(os.path.join(checkpoint_dir, 'model.h5'))
dump(est_c1, os.path.join(checkpoint_dir, 'est_c1.joblib'))
dump(est_c2, os.path.join(checkpoint_dir, 'est_c2.joblib'))
return model, est_c1, est_c2
# get data loader for Y
Y_train_size = int(0.9 * len(Y_dset))
Y_test_size = len(Y_dset) - Y_train_size
Y_train_dset, Y_test_dset = random_split(Y_dset, [Y_train_size, Y_test_size])
Y_train_dataloader = DataLoader(Y_train_dset,
batch_size=batch_size,
shuffle=True,
drop_last=True)
Y_test_dataloader = DataLoader(Y_test_dset,
batch_size=batch_size,
shuffle=True)
torch.save(Y_test_dataloader, os.path.join(save_path, 'Y_test_dataloader.pth'))
# create models
G_XtoY, G_YtoX, D_X, D_Y, device = network.create_model()
# Get generator parameters
g_params = list(G_XtoY.parameters()) + list(G_YtoX.parameters())
# Create optimizers for the generators and discriminators
g_optimizer = optim.Adam(g_params, lr, [beta1, beta2])
d_x_optimizer = optim.Adam(D_X.parameters(), lr, [beta1, beta2])
d_y_optimizer = optim.Adam(D_Y.parameters(), lr, [beta1, beta2])
# scheduler
if is_scheduler:
g_lr_scheduler = torch.optim.lr_scheduler.StepLR(g_optimizer,
step_size=scheduler_step,
gamma=0.1,
last_epoch=-1)