def main():
print("Start... => main.py PID: %s" % (os.getpid()))
opt = get_opt()
# torch.manual_seed(opt.seed)
# random.seed(opt.seed)
# np.random.seed(opt.seed)
# print(opt.gpus)
# if opt.gpus and torch.cuda.is_available():
# print("opt.gpus: ", opt.gpus)
# gpu_list = [int(k) for k in opt.gpus.split(",")]
# print("gpu_list: ", gpu_list)
# print("gpu_list[0]: ", gpu_list[0])
# print("type(gpu_list[0]): ", type(gpu_list[0]))
# os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
# all_dict = load_dict()
if opt.train_mode == "train":
data_loader, all_dict = load_data(opt)
print('Loaded dataset sucessfully.')
model = create_code_retrieval_model(opt, all_dict)
print('created model..')
cos_ranking_loss = CosRankingLoss(
opt).cuda() if torch.cuda.is_available() else CosRankingLoss(opt)
optim = torch.optim.Adam(model.parameters(), opt.lr)
print('created criterion and optim') # standard
trainer = CodeRetrievalTrainer(model, data_loader, all_dict, opt)
# begin to training !!!
trainer.train(cos_ranking_loss, optim, opt.train_epoch)
if opt.train_mode == "query":
data_loader, all_dict, code_body, func_name = load_data(
opt, model='n_query')
model = create_code_retrieval_model(opt, all_dict)
model.load_state_dict(
torch.load('./model/Models/model_0.pt',
map_location=lambda storage, loc: storage))
print('Loaded dataset sucessfully.')
code_searcher = CodereRerievalSearcher(10, code_body, func_name,
all_dict, opt, model,
data_loader)
result = code_searcher.search('set working directory')
code_searcher.save_result()
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torch.autograd import Variable
from tensorboardX import SummaryWriter
import models.dcgan_model as dcgan
from utils import save_checkpoint, load_checkpoint, save_checkpoint_wo_step
from utils import board_add_image, board_add_images
from utils import save_image_historys_gif
from utils import inception_score
from opt import get_opt
opt = get_opt()
print(opt)
# =======
# device
# =======
cuda = True if torch.cuda.is_available() else False
if (opt.device == "gpu"):
if (cuda):
device = torch.device("cuda")
else:
device = torch.device("cpu")
else:
device = torch.device("cpu")
# ===================
def main():
print("Start... => main.py PID: %s" % (os.getpid()))
opt = get_opt()
torch.manual_seed(opt.seed)
random.seed(opt.seed)
np.random.seed(opt.seed)
if opt.gpus:
print("opt.gpus: ", opt.gpus)
gpu_list = [int(k) for k in opt.gpus.split(",")]
print("gpu_list: ", gpu_list)
print("gpu_list[0]: ", gpu_list[0])
print("type(gpu_list[0]): ", type(gpu_list[0]))
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
all_dict = load_dict(opt)
if opt.train_mode == "train":
if opt.validation_with_metric:
train_dataset, train_dataloader, val_dataset, val_dataloader, \
test_dataset, test_dataloader, query_dataset, query_dataloader = load_data(opt, all_dict)
else:
train_dataset, train_dataloader, val_dataset, val_dataloader = load_data(
opt, all_dict)
elif opt.train_mode == "test":
test_dataset, test_dataloader, query_dataset, query_dataloader = load_data(
opt, all_dict)
print('Loaded dataset sucessfully.')
if opt.train_mode == "train":
model = create_model_code_retrieval(opt, train_dataset, all_dict)
print('created model..')
cos_ranking_loss = CosRankingLoss(opt).cuda()
optim = torch.optim.Adam(model.parameters(), opt.lr)
print('created criterion and optim')
if opt.validation_with_metric:
pretrainer = CodeRetrievalTrainer(model,
train_dataset,
train_dataloader,
val_dataset,
val_dataloader,
all_dict,
opt,
metric_data=[
test_dataset,
test_dataloader,
query_dataset,
query_dataloader
])
else:
pretrainer = CodeRetrievalTrainer(model, train_dataset,
train_dataloader, val_dataset,
val_dataloader, all_dict, opt)
pretrainer.train(cos_ranking_loss, optim, opt.pretrain_sl_epoch)
elif opt.train_mode == "test":
model = create_model_code_retrieval(opt, test_dataset, all_dict)
print('created model..')
evaluator = CodeRetrievalEvaluator(model=model,
dataset_list=[
test_dataset, test_dataloader,
query_dataset, query_dataloader
],
flag_for_val=False,
all_dict=all_dict,
opt=opt)
evaluator.retrieval(pred_file=opt.retrieval_pred_file)
if opt.use_val_as_codebase:
evaluator.eval_retrieval_json_result(
pred_file=opt.retrieval_pred_file)