def main():
from arguments.qgen_args import qgen_arguments
from data_provider.qgen_dataset import prepare_dataset
from process_data.tokenizer import GWTokenizer
parser = qgen_arguments()
args, _ = parser.parse_known_args()
args = vars(args)
tokenizer = GWTokenizer('./../data/dict.json')
loader = prepare_dataset("./../data/", "test", args, tokenizer)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = QGenNetwork(args, tokenizer, device).to(device)
optimizer = torch.optim.Adam(model.parameters(), args["lr"])
data_iter = iter(loader)
model.train()
for i in range(5):
batch = next(data_iter)
optimizer.zero_grad()
model.zero_grad()
cls_loss, gen_loss = model(batch)
if args["task"] == "cls":
loss = cls_loss
elif args["task"] == "gen":
loss = gen_loss
else:
loss = cls_loss + gen_loss
loss.backward()
_ = torch.nn.utils.clip_grad_norm_(model.parameters(), args["clip_val"])
optimizer.step()
print("loss: {:.4f}".format(loss.item()))
model.eval()
batch = next(data_iter)
result, pred = model.generate(batch)
print("generate")
print(pred[0].item())
print(tokenizer.decode(result[0]))
def main(args):
param_file = save_path.format("params.pth.tar")
data_dir = "./../data/"
model_name = args["model"].lower()
tokenizer = GWTokenizer('./../data/dict.json')
if model_name == "cat_base":
from models.qgen.qgen_cat_base import QGenNetwork
from data_provider.qgen_dataset import prepare_dataset
elif model_name == "hrnn":
from models.qgen.qgen_hrnn import QGenNetwork
from data_provider.qgen_dataset import prepare_dataset
elif model_name == "cat_accu":
from models.qgen.qgen_cat_accu import QGenNetwork
from data_provider.qgen_dataset import prepare_dataset
elif model_name == "cat_attn":
from models.qgen.qgen_cat_attn import QGenNetwork
from data_provider.qgen_dataset import prepare_dataset
else:
print(model_name)
from models.qgen.qgen_baseline import QGenNetwork
from data_provider.qgen_baseline_dataset import prepare_dataset
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if args["option"] == "train":
if args["image"] is False and args["object"] is False:
print("default object")
args["object"] = True
args["image_arch"] = "rcnn"
args["image_dim"] = 2048
with open(save_path.format("args.json"), mode="w") as f:
json.dump(args, f, indent=2, ensure_ascii=False)
logger.info(args)
model = QGenNetwork(args, tokenizer, device).to(device)
train_loader, val_loader = prepare_dataset(data_dir, "train", args,
tokenizer)
train(model, args, train_loader, val_loader, param_file)
else:
with open(save_path.format("args.json"), mode="r") as f:
saved_args = json.load(f)
saved_args["option"] = "test"
args = saved_args
logger.info(args)
model = QGenNetwork(args, tokenizer, device).to(device)
testloader = prepare_dataset(data_dir, "test", args, tokenizer)
test(model, args, testloader, param_file)