mirror of
https://github.com/hanzoai/jin.git
synced 2026-10-30 18:33:17 +00:00
63 lines
1.4 KiB
YAML
63 lines
1.4 KiB
YAML
# general model params
|
|
model : ijepa
|
|
model_save_name : ijepa
|
|
# good params for jepa
|
|
h : 256
|
|
w : 256
|
|
d : 64 # make this null if you don't want it
|
|
dim : 512
|
|
# good params for autoencoder
|
|
# h : 8
|
|
# w : 8
|
|
# d : 8
|
|
in_channels : 1
|
|
crop_h : 256
|
|
crop_w : 256
|
|
patch_size : [32, 32, 8]
|
|
hw_reduction : 16
|
|
|
|
# these only apply to masked AEs
|
|
start_mask_prob : 0.6
|
|
patcher_type : "conv3d"
|
|
|
|
# these only apply to ijepa
|
|
n_targets : 4 # number of target patches to use
|
|
decorrelation_weight : 0.01
|
|
mask_prob_update_amount : 0.002
|
|
target_scale_fraction_range : [0.2, 0.3]
|
|
|
|
# these only apply to saccade jepa
|
|
affine_cos_loss_weight : 1.0
|
|
affine_mag_loss_weight : 0.5
|
|
cycle_loss_weight : 0.25
|
|
variance_weight : 25
|
|
covariance_weight : 1
|
|
vicreg_gamma : 4
|
|
|
|
# image saving only apply for AEs
|
|
save_images_every : 1
|
|
save_images_every_mini_epoch : 5
|
|
save_mini_epoch : False
|
|
|
|
# general training params
|
|
n_epochs : 75
|
|
mini_epoch_len : 4000 # number of *steps* per mini-epoch
|
|
batch_size : 8
|
|
accumulation_steps : 128
|
|
save_every : 2
|
|
dataset : imagenet
|
|
data_path : ""
|
|
use_util_norm : False
|
|
pretrain_projections : 1
|
|
warmup_epoch_fraction : 100 # fraction of epochs to use for warmup
|
|
|
|
val_every : 0 # run validation tasks every n mini-epochs, 0 means no validation
|
|
tests : ["corr_dim", "plot", "knn"] # the tests to run for validation
|
|
|
|
lr : 0.00005
|
|
min_lr : 0.00001
|
|
ema_decay : 0.9996
|
|
weight_decay : 0.000001
|
|
beta1 : 0.9
|
|
beta2 : 0.99
|
|
clip_grad_norm : 1.0 |