import torch.nn as nn
import numpy as np
from library import inputs
from Utils.checkpoints import save_context, Logger
from Utils import flags
from Utils import config
import Torture
from library import loss_triplegan, evaluation
import library.loss_cla as loss_classifier
from library.mean_teacher import optim_weight_swa
FLAGS = flags.FLAGS
KEY_ARGUMENTS = config.load_config(FLAGS.config_file)
text_logger, MODELS_FOLDER, SUMMARIES_FOLDER = save_context(
__file__, KEY_ARGUMENTS)
# FLAGS.g_model_name = FLAGS.model_name
# FLAGS.d_model_name = FLAGS.model_name
torch.manual_seed(1234)
torch.cuda.manual_seed(1235)
np.random.seed(1236)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
FLAGS.device = device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
n_iter_d = 5 if "sngan" in FLAGS.g_model_name else 1
def sigmoid_rampup(global_step, start_iter, end_iter):
flags.DEFINE_boolean("-combine_with_JPEG", default=False)
flags.DEFINE_boolean("-combine_with_Grayscale", default=False)
flags.DEFINE_argument("-Raffprobability", type=float, default=0.2)
flags.DEFINE_argument("-JPEGquality", type=int, default=75)
flags.DEFINE_argument("-eps", type=float, default=8.)
flags.DEFINE_argument("-numtest", type=int, default=50000)
FLAGS = flags.FLAGS
logger, MODELS_FOLDER, SUMMARIES_FOLDER = save_context(__file__, FLAGS, config)
logger.info("build dataloader")
num_test = FLAGS.numtest
num_pool = 10000
# num of times to sample for MI
num_sample = FLAGS.num_sample
def onehot(ind, num_cla):
vector = np.zeros([num_cla])
vector[ind] = 1
return vector.astype(np.float32)
eps_ = 2.
clip_min, clip_max = -1., 1.
from Torture.utils import distributions
from Torture import shortcuts
from ESM.models import Res18_Quadratic
from ESM.data import inf_train_gen_cifar
FILES_TO_BE_SAVED = ["./", './configs', './ESM']
KEY_ARGUMENTS = ['data', 'esm_eps', 'esm_type']
CONFIG = {
"FILES_TO_BE_SAVED": FILES_TO_BE_SAVED,
"KEY_ARGUMENTS": KEY_ARGUMENTS
}
FLAGS = flags.FLAGS
config.load_config(FLAGS.config_file)
text_logger, MODELS_FOLDER, SUMMARIES_FOLDER = save_context(
__file__, FLAGS, CONFIG)
shutil.copy(FLAGS.config_file, os.path.join(SUMMARIES_FOLDER, "config.yaml"))
torch.manual_seed(1234)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(12345)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Res18_Quadratic(3, FLAGS.n_chan, 32).to(device)
optimizer = torch.optim.Adam(model.parameters(), 5e-5)
checkpoint_io = Torture.utils.checkpoint.CheckpointIO(
checkpoint_dir=MODELS_FOLDER)
checkpoint_io.register_modules(model=model)
logger = Logger(log_dir=SUMMARIES_FOLDER)
