def evaluate_model(config, model_name):
"""
Evaluation of testset
"""
config_file = Path(config).resolve()
cfg = read_config(config_file)
if not os.path.exists(cfg['project_path'] + 'model/evaluate/'):
os.mkdir(cfg['project_path'] + '/' + 'model/evaluate/')
use_gpu = torch.cuda.is_available()
if use_gpu:
print("Using CUDA")
print('GPU active:', torch.cuda.is_available())
print('GPU used:', torch.cuda.get_device_name(0))
else:
print("CUDA is not working!")
print("\n\nEvaluation of %s model. \n" % model_name)
eval_temporal(cfg, use_gpu, model_name)
print(
"You can find the results of the evaluation in '/Your-VAME-Project-Apr30-2020/model/evaluate/' \n"
"OPTIONS:\n"
"- vame.behavior_segmentation() to identify behavioral motifs.\n"
"- re-run the model for further fine tuning. Check again with vame.evaluate_model()"
)
def create_trainset(config):
config_file = Path(config).resolve()
cfg = read_config(config_file)
legacy = cfg['legacy']
if not os.path.exists(
os.path.join(cfg['project_path'], 'data', 'train', "")):
os.mkdir(os.path.join(cfg['project_path'], 'data', 'train', ""))
files = []
if cfg['all_data'] == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to train on " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
for file in cfg['video_sets']:
files.append(file)
print("Creating training dataset...")
if legacy == False:
traindata(cfg, files, cfg['test_fraction'], cfg['num_features'],
cfg['savgol_filter'])
else:
traindata_legacy(cfg, files, cfg['test_fraction'], cfg['num_features'],
cfg['savgol_filter'])
def create_trainset(config):
config_file = Path(config).resolve()
cfg = read_config(config_file)
legacy = cfg['legacy']
if not os.path.exists(
os.path.join(cfg['project_path'], 'data', 'train', "")):
os.mkdir(os.path.join(cfg['project_path'], 'data', 'train', ""))
files = []
if cfg['all_data'] == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to train on " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
for file in cfg['video_sets']:
files.append(file)
print("Creating training dataset...")
if cfg['robust'] == True:
print("Using robust setting to eliminate outliers! IQR factor: %d" %
cfg['iqr_factor'])
if legacy == False:
traindata(cfg, files, cfg['test_fraction'], cfg['num_features'],
cfg['savgol_filter'])
else:
traindata_legacy(cfg, files, cfg['test_fraction'], cfg['num_features'],
cfg['savgol_filter'])
print(
"A training and test set has been created. Now everything is ready to train a variational autoencoder"
"via vame.train_model() ...")
def community(config, show_umap=False, cut_tree=None):
config_file = Path(config).resolve()
cfg = read_config(config_file)
model_name = cfg['model_name']
n_cluster = cfg['n_cluster']
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to write motif videos for your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
labels = get_labels(cfg, files, model_name, n_cluster)
transition_matrices = compute_transition_matrices(files, labels, n_cluster)
communities_all, trees = create_community_bag(files, labels,
transition_matrices,
cut_tree, n_cluster)
community_labels_all = get_community_labels(files, labels, communities_all)
for idx, file in enumerate(files):
path_to_file = os.path.join(cfg['project_path'], "results", file,
model_name, 'kmeans-' + str(n_cluster), "")
if not os.path.exists(os.path.join(path_to_file, "community")):
os.mkdir(os.path.join(path_to_file, "community"))
np.save(
os.path.join(path_to_file, "community",
"transition_matrix_" + file + '.npy'),
transition_matrices[idx])
np.save(
os.path.join(path_to_file, "community",
"community_label_" + file + '.npy'),
community_labels_all[idx])
with open(
os.path.join(path_to_file, "community",
"hierarchy" + file + ".pkl"),
"wb") as fp: #Pickling
pickle.dump(communities_all[idx], fp)
if show_umap == True:
embed = umap_embedding(cfg, file, model_name, n_cluster)
umap_vis(cfg, files, embed, community_labels_all[idx])
def egocentric_alignment(config,
pose_ref_index=[0, 5],
crop_size=(300, 300),
use_video=False,
video_format='.mp4',
check_video=False):
""" Happy aligning """
#config parameters
config_file = Path(config).resolve()
cfg = read_config(config_file)
path_to_file = cfg['project_path']
filename = cfg['video_sets']
video_format = video_format
crop_size = crop_size
# call function and save into your VAME data folder
for file in filename:
print("Egocentric alignment for file %s" % file)
egocentric_time_series, frames = alignment(path_to_file,
file,
pose_ref_index,
video_format,
crop_size,
use_video=use_video,
check_video=check_video)
np.save(os.path.join(path_to_file, 'data', file, file + '-PE-seq.npy'),
egocentric_time_series)
def evaluate_model(config, model_name, suffix=None):
"""
Evaluation of testset
"""
config_file = Path(config).resolve()
cfg = read_config(config_file)
legacy = cfg['legacy']
model_name = cfg['model_name']
if not os.path.exists(
os.path.join(cfg['project_path'], "model", "evaluate")):
os.mkdir(os.path.join(cfg['project_path'], "model", "evaluate"))
use_gpu = torch.cuda.is_available()
if use_gpu:
print("Using CUDA")
print('GPU active:', torch.cuda.is_available())
print('GPU used:', torch.cuda.get_device_name(0))
else:
torch.device("cpu")
print("CUDA is not working, or a GPU is not found; using CPU!")
print("\n\nEvaluation of %s model. \n" % model_name)
eval_temporal(cfg, use_gpu, model_name, suffix=suffix)
print(
"You can find the results of the evaluation in '/Your-VAME-Project-Apr30-2020/model/evaluate/' \n"
"OPTIONS:\n"
"- vame.behavior_segmentation() to identify behavioral motifs.\n"
"- re-run the model for further fine tuning. Check again with vame.evaluate_model()"
)
def generative_model(config, mode="sampling"):
config_file = Path(config).resolve()
cfg = read_config(config_file)
model_name = cfg['model_name']
n_cluster = cfg['n_cluster']
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to write motif videos for your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
model = load_model(cfg, model_name)
for file in files:
path_to_file = os.path.join(cfg['project_path'], "results", file,
model_name, 'kmeans-' + str(n_cluster), "")
if mode == "sampling":
latent_vector = np.load(
os.path.join(path_to_file, 'latent_vector_' + file + '.npy'))
random_generative_samples(cfg, model, latent_vector)
if mode == "reconstruction":
latent_vector = np.load(
os.path.join(path_to_file, 'latent_vector_' + file + '.npy'))
random_reconstruction_samples(cfg, model, latent_vector)
if mode == "centers":
cluster_center = np.load(
os.path.join(path_to_file, 'cluster_center_' + file + '.npy'))
visualize_cluster_center(cfg, model, cluster_center)
if mode == "motifs":
latent_vector = np.load(
os.path.join(path_to_file, 'latent_vector_' + file + '.npy'))
labels = np.load(
os.path.join(path_to_file, "",
str(n_cluster) + '_km_label_' + file + '.npy'))
random_generative_samples_motif(cfg, model, latent_vector, labels,
n_cluster)
def csv_to_numpy(config, datapath):
"""
This is a function to convert your pose-estimation.csv file to a numpy array.
Note that this code is only useful for data which is a priori egocentric, i.e. head-fixed
or otherwise restrained animals.
example use:
vame.csv_to_npy('pathto/your/config/yaml', 'path/toYourFolderwithCSV/')
"""
config_file = Path(config).resolve()
cfg = read_config(config_file)
path_to_file = cfg['project_path']
filename = cfg['video_sets']
for file in filename:
# Read in your .csv file, skip the first two rows and create a numpy array
data = pd.read_csv(datapath + file + '.csv', skiprows=2)
data_mat = pd.DataFrame.to_numpy(data)
data_mat = data_mat[:, 1:]
# get the number of bodyparts, their x,y-position and the confidence from DeepLabCut
bodyparts = int(np.size(data_mat[0, :]) / 3)
positions = []
confidence = []
idx = 0
for i in range(bodyparts):
positions.append(data_mat[:, idx:idx + 2])
confidence.append(data_mat[:, idx + 2])
idx += 3
body_position = np.concatenate(positions, axis=1)
con_arr = np.array(confidence)
# find low confidence and set them to NaN (vame.create_trainset(config) will interpolate these NaNs)
body_position_nan = []
idx = -1
for i in range(bodyparts * 2):
if i % 2 == 0:
idx += 1
seq = body_position[:, i]
seq[con_arr[idx, :] < .99] = np.NaN
body_position_nan.append(seq)
final_positions = np.array(body_position_nan)
# save the final_positions array with np.save()
np.save(os.path.join(path_to_file, 'data', file, file + "-PE-seq.npy"),
final_positions)
print("conversion from DeepLabCut csv to numpy complete...")
print(
"Your data is now ine right format and you can call vame.create_trainset()"
)
def behavior_segmentation(config,
model_name=None,
cluster_method='kmeans',
n_cluster=[30]):
config_file = Path(config).resolve()
cfg = read_config(config_file)
for folders in cfg['video_sets']:
if not os.path.exists(
os.path.join(cfg['project_path'], "results", folders, "",
model_name)):
os.mkdir(
os.path.join(cfg['project_path'], "results", folders, "",
model_name))
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to qunatify your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
use_gpu = torch.cuda.is_available()
if use_gpu:
print("Using CUDA")
print('GPU active:', torch.cuda.is_available())
print('GPU used:', torch.cuda.get_device_name(0))
else:
print("CUDA is not working! Attempting to use the CPU...")
torch.device("cpu")
z, z_logger = temporal_quant(cfg, model_name, files, use_gpu)
cluster_latent_space(cfg, files, z, z_logger, cluster_method, n_cluster,
model_name)
def behavior_quantification(config, model_name, cluster_method='kmeans', n_cluster=30, rename=False):
config_file = Path(config).resolve()
cfg = read_config(config_file)
files = []
if cfg['all_data'] == 'No':
all_flag = input("Do you want to quantify your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
if rename:
count=len(files)
for f in files[:count]:
suffix=f.split('_')[-1]
fn = f.replace(suffix, rename[suffix])
files.append(fn)
files = files[count:]
for file in files:
path_to_file=cfg['project_path']+'results/'+file+'/'+model_name+'/'+cluster_method+'-'+str(n_cluster)
if not os.path.exists(path_to_file+'/behavior_quantification/'):
os.mkdir(path_to_file+'/behavior_quantification/')
if rename:
for filename in glob.iglob(path_to_file+'/'+str(n_cluster) + '*.npy'):
file, ext = os.path.splitext(filename.split('/')[-1])
file = file.split(str(n_cluster)+'_km_label_')[-1]
get_network(path_to_file, file, cluster_method, n_cluster)
def create_trainset(config, check_parameter=False):
config_file = Path(config).resolve()
cfg = read_config(config_file)
legacy = cfg['legacy']
fixed = cfg['egocentric_data']
if not os.path.exists(
os.path.join(cfg['project_path'], 'data', 'train', "")):
os.mkdir(os.path.join(cfg['project_path'], 'data', 'train', ""))
files = []
if cfg['all_data'] == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to train on " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
for file in cfg['video_sets']:
files.append(file)
print("Creating training dataset...")
if cfg['robust'] == True:
print("Using robust setting to eliminate outliers! IQR factor: %d" %
cfg['iqr_factor'])
if fixed == False:
print(
"Creating trainset from the vame.egocentrical_alignment() output ")
traindata_aligned(cfg, files, cfg['test_fraction'],
cfg['num_features'], cfg['savgol_filter'],
check_parameter)
else:
print("Creating trainset from the vame.csv_to_numpy() output ")
traindata_fixed(cfg, files, cfg['test_fraction'], cfg['num_features'],
cfg['savgol_filter'], check_parameter)
if check_parameter == False:
print(
"A training and test set has been created. Next step: vame.train_model()"
)
def egocentric_alignment(config,
pose_ref_index=[0, 5],
crop_size=(300, 300),
use_video=False,
video_format='.mp4',
check_video=False):
""" Happy aligning """
#config parameters
config_file = Path(config).resolve()
cfg = read_config(config_file)
path_to_file = cfg['project_path']
filename = cfg['video_sets']
confidence = cfg['pose_confidence']
video_format = video_format
crop_size = crop_size
if cfg['egocentric_data'] == True:
raise ValueError(
"The config.yaml indicates that the data is not egocentric. Please check the parameter egocentric_data"
)
# call function and save into your VAME data folder
for file in filename:
print("Aligning data %s, Pose confidence value: %.2f" %
(file, confidence))
egocentric_time_series, frames = alignment(path_to_file,
file,
pose_ref_index,
video_format,
crop_size,
confidence,
use_video=use_video,
check_video=check_video)
np.save(os.path.join(path_to_file, 'data', file, file + '-PE-seq.npy'),
egocentric_time_series)
# np.save(os.path.join(path_to_file,'data/',file,"",file+'-PE-seq.npy', egocentric_time_series))
print(
"Your data is now ine right format and you can call vame.create_trainset()"
)
def behavior_quantification(config,
model_name,
cluster_method='kmeans',
n_cluster=30):
config_file = Path(config).resolve()
cfg = read_config(config_file)
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to quantify your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
for file in files:
print('Processing:' + file)
path_to_file = cfg[
'project_path'] + 'results/' + file + '/' + model_name + '/' + cluster_method + '-' + str(
n_cluster)
if not os.path.exists(path_to_file + '/behavior_quantification/'):
os.mkdir(path_to_file + '/behavior_quantification/')
get_network(path_to_file, file, cluster_method, n_cluster)
def community_videos(config, videoType='.mp4'):
config_file = Path(config).resolve()
cfg = read_config(config_file)
model_name = cfg['model_name']
n_cluster = cfg['n_cluster']
flag = 'community'
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to write motif videos for your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
print("Cluster size is: %d " % n_cluster)
for file in files:
path_to_file = os.path.join(cfg['project_path'], "results", file,
model_name, 'kmeans-' + str(n_cluster), "")
if not os.path.exists(os.path.join(path_to_file, "community_videos")):
os.mkdir(os.path.join(path_to_file, "community_videos"))
get_cluster_vid(cfg, path_to_file, file, n_cluster, videoType, flag)
print("All videos have been created!")
def motif_videos(config, model_name, cluster_method="kmeans", n_cluster=[30]):
config_file = Path(config).resolve()
cfg = read_config(config_file)
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to write motif videos for your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
for cluster in n_cluster:
print("Cluster size %d " % cluster)
for file in files:
path_to_file = cfg[
'project_path'] + 'results/' + file + '/' + model_name + '/' + cluster_method + '-' + str(
cluster)
if not os.path.exists(path_to_file + '/cluster_videos/'):
os.mkdir(path_to_file + '/cluster_videos/')
get_cluster_vid(cfg, path_to_file, file, cluster)
def create_trainset(config):
config_file = Path(config).resolve()
cfg = read_config(config_file)
path_to_file = cfg['project_path'] + 'data/'
if not os.path.exists(cfg['project_path'] + 'data/train/'):
os.mkdir(path_to_file + 'train/')
files = []
if cfg['all_data'] == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to train on " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
for file in cfg['video_sets']:
files.append(file)
print("Creating training dataset.")
temporal_traindata(cfg, files, cfg['test_fraction'], cfg['num_features'],
cfg['savgol_filter'])
def gif(config,
pose_ref_index,
subtract_background=True,
start=None,
length=500,
max_lag=30,
label='community',
file_format='.mp4',
crop_size=(300, 300)):
config_file = Path(config).resolve()
cfg = read_config(config_file)
model_name = cfg['model_name']
n_cluster = cfg['n_cluster']
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to write motif videos for your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
for file in files:
path_to_file = os.path.join(cfg['project_path'], "results", file,
model_name, 'kmeans-' + str(n_cluster), "")
if not os.path.exists(os.path.join(path_to_file, "gif_frames")):
os.mkdir(os.path.join(path_to_file, "gif_frames"))
embed = np.load(
os.path.join(path_to_file, "community",
"umap_embedding_" + file + '.npy'))
try:
embed = np.load(
os.path.join(path_to_file, "", "community", "",
"umap_embedding_" + file + ".npy"))
num_points = cfg['num_points']
if num_points > embed.shape[0]:
num_points = embed.shape[0]
except:
print("Compute embedding for file %s" % file)
reducer = umap.UMAP(n_components=2,
min_dist=cfg['min_dist'],
n_neighbors=cfg['n_neighbors'],
random_state=cfg['random_state'])
latent_vector = np.load(
os.path.join(path_to_file, "",
'latent_vector_' + file + '.npy'))
num_points = cfg['num_points']
if num_points > latent_vector.shape[0]:
num_points = latent_vector.shape[0]
print("Embedding %d data points.." % num_points)
embed = reducer.fit_transform(latent_vector[:num_points, :])
np.save(
os.path.join(path_to_file, "community",
"umap_embedding_" + file + '.npy'), embed)
if label == "motif":
umap_label = np.load(
os.path.join(path_to_file,
str(n_cluster) + "_km_label_" + file + '.npy'))
elif label == "community":
umap_label = np.load(
os.path.join(path_to_file, "community",
"community_label_" + file + '.npy'))
elif label == None:
umap_label = None
if start == None:
start = np.random.choice(embed[:num_points].shape[0] - length)
else:
start = start
frames = get_animal_frames(cfg, file, pose_ref_index, start, length,
subtract_background, file_format, crop_size)
create_video(path_to_file, file, embed, umap_label, frames, start,
length, max_lag, num_points)
def rnn_model(config,
model_name,
pretrained_weights=False,
pretrained_model=None):
config_file = Path(config).resolve()
cfg = read_config(config_file)
print("Train RNN model!")
if not os.path.exists(cfg['project_path'] + '/' + 'model/best_model'):
os.mkdir(cfg['project_path'] + '/model/' + 'best_model')
os.mkdir(cfg['project_path'] + '/model/' + 'best_model/snapshots')
os.mkdir(cfg['project_path'] + '/model/' + 'model_losses')
# make sure torch uses cuda for GPU computing
use_gpu = torch.cuda.is_available()
if use_gpu:
print("Using CUDA")
print('GPU active:', torch.cuda.is_available())
print('GPU used:', torch.cuda.get_device_name(0))
else:
print("CUDA is not working!")
raise NotImplementedError('GPU Computing is required!')
""" HYPERPARAMTERS """
# General
CUDA = use_gpu
SEED = 19
TRAIN_BATCH_SIZE = cfg['batch_size']
TEST_BATCH_SIZE = int(cfg['batch_size'] / 4)
EPOCHS = cfg['max_epochs']
ZDIMS = cfg['zdims']
BETA = cfg['beta']
SNAPSHOT = cfg['model_snapshot']
LEARNING_RATE = cfg['learning_rate']
NUM_FEATURES = cfg['num_features']
TEMPORAL_WINDOW = cfg['time_window'] * 2
FUTURE_DECODER = cfg['prediction_decoder']
FUTURE_STEPS = cfg['prediction_steps']
# RNN
hidden_size_layer_1 = cfg['hidden_size_layer_1']
hidden_size_layer_2 = cfg['hidden_size_layer_2']
hidden_size_rec = cfg['hidden_size_rec']
hidden_size_pred = cfg['hidden_size_pred']
dropout_encoder = cfg['dropout_encoder']
dropout_rec = cfg['dropout_rec']
dropout_pred = cfg['dropout_pred']
# Loss
MSE_REC_REDUCTION = cfg['mse_reconstruction_reduction']
MSE_PRED_REDUCTION = cfg['mse_prediction_reduction']
KMEANS_LOSS = cfg['kmeans_loss']
KMEANS_LAMBDA = cfg['kmeans_lambda']
KL_START = cfg['kl_start']
ANNEALTIME = cfg['annealtime']
anneal_function = cfg['anneal_function']
optimizer_scheduler = cfg['scheduler']
BEST_LOSS = 999999
convergence = 0
print('Latent Dimensions: %d, Beta: %d, lr: %.4f' %
(ZDIMS, BETA, LEARNING_RATE))
# simple logging of diverse losses
train_losses = []
test_losses = []
kmeans_losses = []
kl_losses = []
weight_values = []
mse_losses = []
fut_losses = []
torch.manual_seed(SEED)
if CUDA:
torch.cuda.manual_seed(SEED)
model = RNN_VAE(TEMPORAL_WINDOW, ZDIMS, NUM_FEATURES, FUTURE_DECODER,
FUTURE_STEPS, hidden_size_layer_1, hidden_size_layer_2,
hidden_size_rec, hidden_size_pred, dropout_encoder,
dropout_rec, dropout_pred).cuda()
if pretrained_weights:
if os.path.exists(cfg['project_path'] + '/' + 'model/' +
'pretrained_model/' + pretrained_model + '.pkl'):
print("Loading pretrained Model: %s" % pretrained_model)
model.load_state_dict(
torch.load(cfg['project_path'] + '/' + 'model/' +
'pretrained_model/' + pretrained_model + '.pkl'),
strict=False)
""" DATASET """
trainset = SEQUENCE_DATASET(cfg['project_path'] + 'data/train/',
data='train_seq.npy',
train=True,
temporal_window=TEMPORAL_WINDOW)
testset = SEQUENCE_DATASET(cfg['project_path'] + 'data/train/',
data='test_seq.npy',
train=False,
temporal_window=TEMPORAL_WINDOW)
train_loader = Data.DataLoader(trainset,
batch_size=TRAIN_BATCH_SIZE,
shuffle=True,
drop_last=True)
test_loader = Data.DataLoader(testset,
batch_size=TEST_BATCH_SIZE,
shuffle=True,
drop_last=True)
optimizer = torch.optim.Adam(model.parameters(),
lr=LEARNING_RATE,
amsgrad=True)
if optimizer_scheduler:
scheduler = StepLR(optimizer, step_size=100, gamma=0.2, last_epoch=-1)
else:
scheduler = StepLR(optimizer, step_size=100, gamma=0, last_epoch=-1)
for epoch in range(1, EPOCHS):
print('Epoch: %d' % epoch)
print('Train: ')
weight, train_loss, km_loss, kl_loss, mse_loss, fut_loss = train(
train_loader, epoch, model, optimizer, anneal_function, BETA,
KL_START, ANNEALTIME, TEMPORAL_WINDOW, FUTURE_DECODER,
FUTURE_STEPS, scheduler, MSE_REC_REDUCTION, MSE_PRED_REDUCTION,
KMEANS_LOSS, KMEANS_LAMBDA, TRAIN_BATCH_SIZE)
print('Test: ')
current_loss, test_loss, test_list = test(
test_loader, epoch, model, optimizer, BETA, weight,
TEMPORAL_WINDOW, MSE_REC_REDUCTION, KMEANS_LOSS, KMEANS_LAMBDA,
FUTURE_DECODER, TEST_BATCH_SIZE)
for param_group in optimizer.param_groups:
print('lr: {}'.format(param_group['lr']))
# logging losses
train_losses.append(train_loss)
test_losses.append(test_loss)
kmeans_losses.append(km_loss)
kl_losses.append(kl_loss)
weight_values.append(weight)
mse_losses.append(mse_loss)
fut_losses.append(fut_loss)
# save best model
if weight > 0.99 and current_loss <= BEST_LOSS:
BEST_LOSS = current_loss
print("Saving model!\n")
torch.save(
model.state_dict(),
cfg['project_path'] + '/' + 'model/' + 'best_model' + '/' +
model_name + '_' + cfg['Project'] + '.pkl')
convergence = 0
else:
convergence += 1
if epoch % SNAPSHOT == 0:
print("Saving model snapshot!\n")
torch.save(
model.state_dict(), cfg['project_path'] + '/' + 'model/' +
'best_model' + '/snapshots/' + model_name + '_' +
cfg['Project'] + '_epoch_' + str(epoch) + '.pkl')
if convergence > cfg['model_convergence']:
print(
'Model converged. Please check your model with vame.evaluate_model(). \n'
'You can also re-run vame.rnn_model() to further improve your model. \n'
'Hint: Set "model_convergence" in your config.yaml to a higher value. \n'
'\n'
'Next: \n'
'Use vame.behavior_segmentation() to identify behavioral motifs in your dataset!'
)
#return
break
# save logged losses
np.save(
cfg['project_path'] + '/' + 'model/' + 'model_losses' +
'/train_losses_' + model_name, train_losses)
np.save(
cfg['project_path'] + '/' + 'model/' + 'model_losses' +
'/test_losses_' + model_name, test_losses)
np.save(
cfg['project_path'] + '/' + 'model/' + 'model_losses' +
'/kmeans_losses_' + model_name, kmeans_losses)
np.save(
cfg['project_path'] + '/' + 'model/' + 'model_losses' +
'/kl_losses_' + model_name, kl_losses)
np.save(
cfg['project_path'] + '/' + 'model/' + 'model_losses' +
'/weight_values_' + model_name, weight_values)
np.save(
cfg['project_path'] + '/' + 'model/' + 'model_losses' +
'/mse_train_losses_' + model_name, mse_losses)
np.save(
cfg['project_path'] + '/' + 'model/' + 'model_losses' +
'/mse_test_losses_' + model_name, current_loss)
np.save(
cfg['project_path'] + '/' + 'model/' + 'model_losses' +
'/fut_losses_' + model_name, fut_losses)
if convergence < cfg['model_convergence']:
print(
'Model seemed to have not reached convergence. You may want to check your model \n'
'with vame.evaluate_model(). If your satisfied you can continue with \n'
'Use vame.behavior_segmentation() to identify behavioral motifs!\n\n'
'OPTIONAL: You can re-run vame.rnn_model() to improve performance.'
)
def csv_to_numpy(config):
"""
This is a function to convert your pose-estimation.csv file to a numpy array.
Note that this code is only useful for data which is a priori egocentric, i.e. head-fixed
or otherwise restrained animals.
example use:
vame.csv_to_npy('pathto/your/config/yaml', 'path/toYourFolderwithCSV/')
"""
config_file = Path(config).resolve()
cfg = read_config(config_file)
path_to_file = cfg['project_path']
filename = cfg['video_sets']
confidence = cfg['pose_confidence']
if cfg['egocentric_data'] == False:
raise ValueError(
"The config.yaml indicates that the data is not egocentric. Please check the parameter egocentric_data"
)
for file in filename:
print(file)
# Read in your .csv file, skip the first two rows and create a numpy array
data = pd.read_csv(os.path.join(path_to_file, "videos",
"pose_estimation", file + '.csv'),
skiprows=3,
header=None)
data_mat = pd.DataFrame.to_numpy(data)
data_mat = data_mat[:, 1:]
pose_list = []
# get the number of bodyparts, their x,y-position and the confidence from DeepLabCut
for i in range(int(data_mat.shape[1] / 3)):
pose_list.append(data_mat[:, i * 3:(i + 1) * 3])
# find low confidence and set them to NaN
for i in pose_list:
for j in i:
if j[2] <= confidence:
j[0], j[1] = np.nan, np.nan
# interpolate NaNs
for i in pose_list:
i = interpol(i)
positions = np.concatenate(pose_list, axis=1)
final_positions = np.zeros(
(data_mat.shape[0], int(data_mat.shape[1] / 3) * 2))
jdx = 0
idx = 0
for i in range(int(data_mat.shape[1] / 3)):
final_positions[:, idx:idx + 2] = positions[:, jdx:jdx + 2]
jdx += 3
idx += 2
# save the final_positions array with np.save()
np.save(os.path.join(path_to_file, 'data', file, file + "-PE-seq.npy"),
final_positions.T)
print("conversion from DeepLabCut csv to numpy complete...")
print(
"Your data is now ine right format and you can call vame.create_trainset()"
)
def motif_videos(config,
model_name,
cluster_method="kmeans",
n_cluster=[30],
rename=None):
config_file = Path(config).resolve()
cfg = read_config(config_file)
model_name = cfg['model_name']
n_cluster = cfg['n_cluster']
flag = 'motif'
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to write motif videos for your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
for cluster in n_cluster:
print("Cluster size %d " % cluster)
if rename:
for file in files:
suffix = file.split('_')[-1]
file = file.replace(suffix, rename[suffix])
path_to_file = cfg[
'project_path'] + 'results/' + file + '/' + model_name + '/' + cluster_method + '-' + str(
cluster)
if not os.path.exists(path_to_file + '/cluster_videos/'):
os.mkdir(path_to_file + '/cluster_videos/')
get_cluster_vid(cfg,
path_to_file,
file,
n_cluster,
videoType,
flag,
cluster_method=cluster_method,
rename=rename)
elif not rename:
for file in files:
path_to_file = cfg[
'project_path'] + 'results/' + file + '/' + model_name + '/' + cluster_method + '-' + str(
cluster)
if not os.path.exists(path_to_file + '/cluster_videos/'):
os.mkdir(path_to_file + '/cluster_videos/')
get_cluster_vid(cfg,
path_to_file,
file,
n_cluster,
videoType,
flag,
cluster_method=cluster_method)
print("All videos have been created!")
def pose_segmentation(config):
config_file = Path(config).resolve()
cfg = read_config(config_file)
legacy = cfg['legacy']
model_name = cfg['model_name']
n_cluster = cfg['n_cluster']
print('Pose segmentation for VAME model: %s \n' % model_name)
if legacy == True:
from segment_behavior import behavior_segmentation
behavior_segmentation(config,
model_name=model_name,
cluster_method='kmeans',
n_cluster=n_cluster)
else:
ind_param = cfg['individual_parameterization']
for folders in cfg['video_sets']:
if not os.path.exists(
os.path.join(cfg['project_path'], "results", folders,
model_name, "")):
os.mkdir(
os.path.join(cfg['project_path'], "results", folders,
model_name, ""))
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to qunatify your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file +
"? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
use_gpu = torch.cuda.is_available()
if use_gpu:
print("Using CUDA")
print('GPU active:', torch.cuda.is_available())
print('GPU used:', torch.cuda.get_device_name(0))
else:
print("CUDA is not working! Attempting to use the CPU...")
torch.device("cpu")
if not os.path.exists(
os.path.join(cfg['project_path'], "results", file, model_name,
'kmeans-' + str(n_cluster), "")):
new = True
# print("Hello1")
model = load_model(cfg, model_name, legacy)
latent_vectors = embedd_latent_vectors(cfg, files, model, legacy)
if ind_param == False:
print(
"For all animals the same k-Means parameterization of latent vectors is applied for %d cluster"
% n_cluster)
labels, cluster_center, motif_usages = same_parameterization(
cfg, files, latent_vectors, n_cluster)
else:
print(
"Individual k-Means parameterization of latent vectors for %d cluster"
% n_cluster)
labels, cluster_center, motif_usages = individual_parameterization(
cfg, files, latent_vectors, n_cluster)
else:
print('\n'
'For model %s a latent vector embedding already exists. \n'
'Parameterization of latent vector with %d k-Means cluster' %
(model_name, n_cluster))
if os.path.exists(
os.path.join(cfg['project_path'], "results", file,
model_name, 'kmeans-' + str(n_cluster), "")):
flag = input(
'WARNING: A parameterization for the chosen cluster size of the model already exists! \n'
'Do you want to continue? A new k-Means assignment will be computed! (yes/no) '
)
else:
flag = 'yes'
if flag == 'yes':
new = True
latent_vectors = []
for file in files:
path_to_latent_vector = os.path.join(
cfg['project_path'], "results", file, model_name,
'kmeans-' + str(n_cluster), "")
latent_vector = np.load(
os.path.join(path_to_latent_vector,
'latent_vector_' + file + '.npy'))
latent_vectors.append(latent_vector)
if ind_param == False:
print(
"For all animals the same k-Means parameterization of latent vectors is applied for %d cluster"
% n_cluster)
labels, cluster_center, motif_usages = same_parameterization(
cfg, files, latent_vectors, n_cluster)
else:
print(
"Individual k-Means parameterization of latent vectors for %d cluster"
% n_cluster)
labels, cluster_center, motif_usages = individual_parameterization(
cfg, files, latent_vectors, n_cluster)
else:
print('No new parameterization has been calculated.')
new = False
# print("Hello2")
if new == True:
for idx, file in enumerate(files):
print(
os.path.join(cfg['project_path'], "results", file, "",
model_name, 'kmeans-' + str(n_cluster), ""))
if not os.path.exists(
os.path.join(cfg['project_path'], "results", file,
model_name, 'kmeans-' + str(n_cluster),
"")):
try:
os.mkdir(
os.path.join(cfg['project_path'], "results", file,
"", model_name,
'kmeans-' + str(n_cluster), ""))
except OSError as error:
print(error)
save_data = os.path.join(cfg['project_path'], "results", file,
model_name,
'kmeans-' + str(n_cluster), "")
np.save(
os.path.join(save_data,
str(n_cluster) + '_km_label_' + file),
labels[idx])
np.save(os.path.join(save_data, 'cluster_center_' + file),
cluster_center[idx])
np.save(os.path.join(save_data, 'latent_vector_' + file),
latent_vectors[idx])
np.save(os.path.join(save_data, 'motif_usage_' + file),
motif_usages[idx])
print(
"You succesfully extracted motifs with VAME! From here, you can proceed running vame.motif_videos() "
"to get an idea of the behavior captured by VAME. This will leave you with short snippets of certain movements."
"To get the full picture of the spatiotemporal dynamic we recommend applying our community approach afterwards."
)
def visualization(config, label=None):
config_file = Path(config).resolve()
cfg = read_config(config_file)
model_name = cfg['model_name']
n_cluster = cfg['n_cluster']
files = []
if cfg['all_data'] == 'No':
all_flag = input(
"Do you want to write motif videos for your entire dataset? \n"
"If you only want to use a specific dataset type filename: \n"
"yes/no/filename ")
else:
all_flag = 'yes'
if all_flag == 'yes' or all_flag == 'Yes':
for file in cfg['video_sets']:
files.append(file)
elif all_flag == 'no' or all_flag == 'No':
for file in cfg['video_sets']:
use_file = input("Do you want to quantify " + file + "? yes/no: ")
if use_file == 'yes':
files.append(file)
if use_file == 'no':
continue
else:
files.append(all_flag)
for idx, file in enumerate(files):
path_to_file = os.path.join(cfg['project_path'], "results", file, "",
model_name, "", 'kmeans-' + str(n_cluster))
try:
embed = np.load(
os.path.join(path_to_file, "", "community", "",
"umap_embedding_" + file + ".npy"))
num_points = cfg['num_points']
if num_points > embed.shape[0]:
num_points = embed.shape[0]
except:
if not os.path.exists(os.path.join(path_to_file, "community")):
os.mkdir(os.path.join(path_to_file, "community"))
print("Compute embedding for file %s" % file)
reducer = umap.UMAP(n_components=2,
min_dist=cfg['min_dist'],
n_neighbors=cfg['n_neighbors'],
random_state=cfg['random_state'])
latent_vector = np.load(
os.path.join(path_to_file, "",
'latent_vector_' + file + '.npy'))
num_points = cfg['num_points']
if num_points > latent_vector.shape[0]:
num_points = latent_vector.shape[0]
print("Embedding %d data points.." % num_points)
embed = reducer.fit_transform(latent_vector[:num_points, :])
np.save(
os.path.join(path_to_file, "community",
"umap_embedding_" + file + '.npy'), embed)
print("Visualizing %d data points.. " % num_points)
if label == None:
umap_vis(file, embed, num_points)
if label == 'motif':
label = np.load(
os.path.join(path_to_file, "",
str(n_cluster) + '_km_label_' + file + '.npy'))
umap_label_vis(file, embed, label, n_cluster, num_points)
if label == "community":
community_label = np.load(
os.path.join(path_to_file, "", "community", "",
"community_label_" + file + ".npy"))
umap_vis_comm(file, embed, community_label, num_points)
def train_model(config):
config_file = Path(config).resolve()
cfg = read_config(config_file)
legacy = cfg['legacy']
model_name = cfg['model_name']
pretrained_weights = cfg['pretrained_weights']
pretrained_model = cfg['pretrained_model']
fixed = cfg['egocentric_data']
print("Train Variational Autoencoder - model name: %s \n" % model_name)
if not os.path.exists(
os.path.join(cfg['project_path'], 'model', 'best_model', "")):
os.mkdir(os.path.join(cfg['project_path'], 'model', 'best_model', ""))
os.mkdir(
os.path.join(cfg['project_path'], 'model', 'best_model',
'snapshots', ""))
os.mkdir(os.path.join(cfg['project_path'], 'model', 'model_losses',
""))
# make sure torch uses cuda for GPU computing
use_gpu = torch.cuda.is_available()
if use_gpu:
print("Using CUDA")
print('GPU active:', torch.cuda.is_available())
print('GPU used: ', torch.cuda.get_device_name(0))
else:
torch.device("cpu")
print("warning, a GPU was not found... proceeding with CPU (slow!) \n")
#raise NotImplementedError('GPU Computing is required!')
""" HYPERPARAMTERS """
# General
CUDA = use_gpu
SEED = 19
TRAIN_BATCH_SIZE = cfg['batch_size']
TEST_BATCH_SIZE = int(cfg['batch_size'] / 4)
EPOCHS = cfg['max_epochs']
ZDIMS = cfg['zdims']
BETA = cfg['beta']
SNAPSHOT = cfg['model_snapshot']
LEARNING_RATE = cfg['learning_rate']
NUM_FEATURES = cfg['num_features']
if fixed == False:
NUM_FEATURES = NUM_FEATURES - 2
TEMPORAL_WINDOW = cfg['time_window'] * 2
FUTURE_DECODER = cfg['prediction_decoder']
FUTURE_STEPS = cfg['prediction_steps']
# RNN
hidden_size_layer_1 = cfg['hidden_size_layer_1']
hidden_size_layer_2 = cfg['hidden_size_layer_2']
hidden_size_rec = cfg['hidden_size_rec']
hidden_size_pred = cfg['hidden_size_pred']
dropout_encoder = cfg['dropout_encoder']
dropout_rec = cfg['dropout_rec']
dropout_pred = cfg['dropout_pred']
noise = cfg['noise']
scheduler_step_size = cfg['scheduler_step_size']
softplus = cfg['softplus']
# Loss
MSE_REC_REDUCTION = cfg['mse_reconstruction_reduction']
MSE_PRED_REDUCTION = cfg['mse_prediction_reduction']
KMEANS_LOSS = cfg['kmeans_loss']
KMEANS_LAMBDA = cfg['kmeans_lambda']
KL_START = cfg['kl_start']
ANNEALTIME = cfg['annealtime']
anneal_function = cfg['anneal_function']
optimizer_scheduler = cfg['scheduler']
BEST_LOSS = 999999
convergence = 0
print(
'Latent Dimensions: %d, Time window: %d, Batch Size: %d, Beta: %d, lr: %.4f\n'
% (ZDIMS, cfg['time_window'], TRAIN_BATCH_SIZE, BETA, LEARNING_RATE))
# simple logging of diverse losses
train_losses = []
test_losses = []
kmeans_losses = []
kl_losses = []
weight_values = []
mse_losses = []
fut_losses = []
torch.manual_seed(SEED)
if legacy == False:
RNN = RNN_VAE
else:
RNN = RNN_VAE_LEGACY
if CUDA:
torch.cuda.manual_seed(SEED)
model = RNN(TEMPORAL_WINDOW, ZDIMS, NUM_FEATURES, FUTURE_DECODER,
FUTURE_STEPS, hidden_size_layer_1, hidden_size_layer_2,
hidden_size_rec, hidden_size_pred, dropout_encoder,
dropout_rec, dropout_pred, softplus).cuda()
else: #cpu support ...
torch.cuda.manual_seed(SEED)
model = RNN(TEMPORAL_WINDOW, ZDIMS, NUM_FEATURES, FUTURE_DECODER,
FUTURE_STEPS, hidden_size_layer_1, hidden_size_layer_2,
hidden_size_rec, hidden_size_pred, dropout_encoder,
dropout_rec, dropout_pred, softplus).to()
if pretrained_weights:
if os.path.exists(
os.path.join(cfg['project_path'], 'model', 'best_model',
pretrained_model + '_' + cfg['Project'] +
'.pkl')):
print("Loading pretrained weights from model: %s\n" %
pretrained_model)
model.load_state_dict(
torch.load(
os.path.join(
cfg['project_path'], 'model', 'best_model',
pretrained_model + '_' + cfg['Project'] + '.pkl')))
KL_START = 0
ANNEALTIME = 1
""" DATASET """
trainset = SEQUENCE_DATASET(os.path.join(cfg['project_path'], "data",
"train", ""),
data='train_seq.npy',
train=True,
temporal_window=TEMPORAL_WINDOW)
testset = SEQUENCE_DATASET(os.path.join(cfg['project_path'], "data",
"train", ""),
data='test_seq.npy',
train=False,
temporal_window=TEMPORAL_WINDOW)
train_loader = Data.DataLoader(trainset,
batch_size=TRAIN_BATCH_SIZE,
shuffle=True,
drop_last=True)
test_loader = Data.DataLoader(testset,
batch_size=TEST_BATCH_SIZE,
shuffle=True,
drop_last=True)
optimizer = torch.optim.Adam(model.parameters(),
lr=LEARNING_RATE,
amsgrad=True)
if optimizer_scheduler:
print('Scheduler step size: %d, Scheduler gamma: %.2f\n' %
(scheduler_step_size, cfg['scheduler_gamma']))
# Thanks to @alexcwsmith for the optimized scheduler contribution
scheduler = ReduceLROnPlateau(optimizer,
'min',
factor=cfg['scheduler_gamma'],
patience=cfg['scheduler_step_size'],
threshold=1e-3,
threshold_mode='rel',
verbose=True)
else:
scheduler = StepLR(optimizer,
step_size=scheduler_step_size,
gamma=1,
last_epoch=-1)
print("Start training... ")
for epoch in range(1, EPOCHS):
print("Epoch: %d" % epoch)
weight, train_loss, km_loss, kl_loss, mse_loss, fut_loss = train(
train_loader, epoch, model, optimizer, anneal_function, BETA,
KL_START, ANNEALTIME, TEMPORAL_WINDOW, FUTURE_DECODER,
FUTURE_STEPS, scheduler, MSE_REC_REDUCTION, MSE_PRED_REDUCTION,
KMEANS_LOSS, KMEANS_LAMBDA, TRAIN_BATCH_SIZE, noise)
current_loss, test_loss, test_list = test(
test_loader, epoch, model, optimizer, BETA, weight,
TEMPORAL_WINDOW, MSE_REC_REDUCTION, KMEANS_LOSS, KMEANS_LAMBDA,
FUTURE_DECODER, TEST_BATCH_SIZE)
# logging losses
train_losses.append(train_loss)
test_losses.append(test_loss)
kmeans_losses.append(km_loss)
kl_losses.append(kl_loss)
weight_values.append(weight)
mse_losses.append(mse_loss)
fut_losses.append(fut_loss)
# save best model
if weight > 0.99 and current_loss <= BEST_LOSS:
BEST_LOSS = current_loss
print("Saving model!")
if use_gpu:
torch.save(
model.state_dict(),
os.path.join(cfg['project_path'], "model", "best_model",
model_name + '_' + cfg['Project'] + '.pkl'))
else:
torch.save(
model.state_dict(),
os.path.join(cfg['project_path'], "model", "best_model",
model_name + '_' + cfg['Project'] + '.pkl'))
convergence = 0
else:
convergence += 1
if epoch % SNAPSHOT == 0:
print("Saving model snapshot!\n")
torch.save(
model.state_dict(),
os.path.join(
cfg['project_path'], 'model', 'best_model', 'snapshots',
model_name + '_' + cfg['Project'] + '_epoch_' +
str(epoch) + '.pkl'))
if convergence > cfg['model_convergence']:
print('Finished training...')
print(
'Model converged. Please check your model with vame.evaluate_model(). \n'
'You can also re-run vame.trainmodel() to further improve your model. \n'
'Make sure to set _pretrained_weights_ in your config.yaml to "true" \n'
'and plug your current model name into _pretrained_model_. \n'
'Hint: Set "model_convergence" in your config.yaml to a higher value. \n'
'\n'
'Next: \n'
'Use vame.pose_segmentation() to identify behavioral motifs in your dataset!'
)
#return
break
# save logged losses
np.save(
os.path.join(cfg['project_path'], 'model', 'model_losses',
'train_losses_' + model_name), train_losses)
np.save(
os.path.join(cfg['project_path'], 'model', 'model_losses',
'test_losses_' + model_name), test_losses)
np.save(
os.path.join(cfg['project_path'], 'model', 'model_losses',
'kmeans_losses_' + model_name), kmeans_losses)
np.save(
os.path.join(cfg['project_path'], 'model', 'model_losses',
'kl_losses_' + model_name), kl_losses)
np.save(
os.path.join(cfg['project_path'], 'model', 'model_losses',
'weight_values_' + model_name), weight_values)
np.save(
os.path.join(cfg['project_path'], 'model', 'model_losses',
'mse_train_losses_' + model_name), mse_losses)
np.save(
os.path.join(cfg['project_path'], 'model', 'model_losses',
'mse_test_losses_' + model_name), current_loss)
np.save(
os.path.join(cfg['project_path'], 'model', 'model_losses',
'fut_losses_' + model_name), fut_losses)
print("\n")
if convergence < cfg['model_convergence']:
print('Finished training...')
print(
'Model seems to have not reached convergence. You may want to check your model \n'
'with vame.evaluate_model(). If your satisfied you can continue. \n'
'Use vame.pose_segmentation() to identify behavioral motifs! \n'
'OPTIONAL: You can re-run vame.train_model() to improve performance.'
)
