def test(args, ckpt_file): batch_size = args["batch_size"] testloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, collate_fn=generate_batch) predictions, targets = [], [] net = TextSentiment(VOCAB_SIZE, EMBED_DIM, NUN_CLASS).to(device) ckpt = torch.load(os.path.join(args["EXPT_DIR"], ckpt_file)) net.load_state_dict(ckpt["model"]) net.eval() correct, total = 0, 0 with torch.no_grad(): for data in tqdm(testloader, desc="Testing"): text, offsets, cls = data text, offsets, cls = text.to(device), offsets.to(device), cls.to( device) outputs = net(text, offsets) _, predicted = torch.max(outputs.data, 1) predictions.extend(predicted.cpu().numpy().tolist()) targets.extend(cls.cpu().numpy().tolist()) total += cls.size(0) correct += (predicted == cls).sum().item() return {"predictions": predictions, "labels": targets}
def infer(args, unlabeled, ckpt_file): unlabeled = Subset(train_dataset, unlabeled) unlabeled_loader = torch.utils.data.DataLoader( unlabeled, batch_size=args["batch_size"], shuffle=False, num_workers=2, collate_fn=generate_batch) net = TextSentiment(VOCAB_SIZE, EMBED_DIM, NUN_CLASS).to(device) ckpt = torch.load(os.path.join(args["EXPT_DIR"], ckpt_file)) net.load_state_dict(ckpt["model"]) net.eval() correct, total = 0, 0 outputs_fin = {} with torch.no_grad(): for i, data in tqdm(enumerate(unlabeled_loader), desc="Inferring"): text, offsets, cls = data text, offsets, cls = text.to(device), offsets.to(device), cls.to( device) outputs = net(text, offsets) _, predicted = torch.max(outputs.data, 1) total += cls.size(0) correct += (predicted == cls).sum().item() for j in range(len(outputs)): outputs_fin[j] = {} outputs_fin[j]["prediction"] = predicted[j].item() outputs_fin[j]["pre_softmax"] = outputs[j].cpu().numpy() return {"outputs": outputs_fin}
def train(args, labeled, resume_from, ckpt_file): batch_size = args["batch_size"] lr = 4.0 momentum = 0.9 epochs = args["train_epochs"] if not os.path.isdir('./.data'): os.mkdir('./.data') global train_dataset, test_dataset train_dataset, test_dataset = text_classification.DATASETS['AG_NEWS']( root='./.data', ngrams=args["N_GRAMS"], vocab=None) global VOCAB_SIZE, EMBED_DIM, NUN_CLASS VOCAB_SIZE = len(train_dataset.get_vocab()) EMBED_DIM = args["EMBED_DIM"] NUN_CLASS = len(train_dataset.get_labels()) trainloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, collate_fn=generate_batch) net = TextSentiment(VOCAB_SIZE, EMBED_DIM, NUN_CLASS).to(device) criterion = nn.CrossEntropyLoss().to(device) optimizer = optim.SGD(net.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1, gamma=0.9) if resume_from is not None: ckpt = torch.load(os.path.join(args["EXPT_DIR"], resume_from)) net.load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) else: getdatasetstate() net.train() for epoch in tqdm(range(epochs), desc="Training"): running_loss = 0.0 train_acc = 0 for i, data in enumerate(trainloader): text, offsets, cls = data text, offsets, cls = text.to(device), offsets.to(device), cls.to( device) outputs = net(text, offsets) loss = criterion(outputs, cls) optimizer.zero_grad() loss.backward() optimizer.step() train_acc += (outputs.argmax(1) == cls).sum().item() running_loss += loss.item() scheduler.step() print("Finished Training. Saving the model as {}".format(ckpt_file)) print("Training accuracy: {}".format( (train_acc / len(train_dataset) * 100))) ckpt = {"model": net.state_dict(), "optimizer": optimizer.state_dict()} torch.save(ckpt, os.path.join(args["EXPT_DIR"], ckpt_file)) return
import pickle import argparse import torch ag_news_label = {1: "World", 2: "Sports", 3: "Business", 4: "Sci/Tec"} WEIGHT_PATH = "../weights/text_news0.2672930294473966.pth" vocab = pickle.load(open(".data/save_vocab.p", "rb")) device = "cuda" if torch.cuda.is_available() else "cpu" VOCAB_SIZE = 1308844 EMBED_DIM = 32 NUM_CLASS = 4 model = TextSentiment(VOCAB_SIZE, EMBED_DIM, NUM_CLASS) checkpoint = torch.load(WEIGHT_PATH, map_location=torch.device('cpu')) model.load_state_dict(checkpoint) model.to(device) def predict(text, model, vocab, ngrams): tokenizer = get_tokenizer("basic_english") with torch.no_grad(): text = torch.tensor([ vocab[token] for token in ngrams_iterator(tokenizer(text), ngrams) ]) output = model(text, torch.tensor([0])) return output.argmax(1).item() + 1 parser = argparse.ArgumentParser( description='Text_classification With Pytorch')