Spaces:
Runtime error
Runtime error
File size: 4,472 Bytes
7900c16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
"""
This script provides an example to wrap TencentPretrain for image classification inference.
"""
import sys
import os
import torch
import argparse
import collections
import torch.nn as nn
from torchvision import transforms
from torchvision.io import read_image
from torchvision.io.image import ImageReadMode
tencentpretrain_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.append(tencentpretrain_dir)
from tencentpretrain.utils.constants import *
from tencentpretrain.utils import *
from tencentpretrain.utils.config import load_hyperparam
from tencentpretrain.utils.seed import set_seed
from tencentpretrain.model_loader import load_model
from tencentpretrain.opts import infer_opts, tokenizer_opts
from tencentpretrain.utils.misc import ZeroOneNormalize
from finetune.run_classifier import Classifier
def data_loader(args, path):
transform = transforms.Compose([
transforms.Resize((args.image_height, args.image_width)),
ZeroOneNormalize()
])
dataset, columns = [], {}
with open(path, mode="r", encoding="utf-8") as f:
src_batch, seg_batch = [], []
for line_id, line in enumerate(f):
if line_id == 0:
for i, column_name in enumerate(line.rstrip("\r\n").split("\t")):
columns[column_name] = i
continue
line = line.rstrip("\r\n").split("\t")
path = line[columns["path"]]
image = read_image(path, ImageReadMode.RGB)
image = image.to(args.device)
src = transform(image)
seg = [1] * ((src.size()[1] // args.patch_size) * (src.size()[2] // args.patch_size) + 1)
src_batch.append(src)
seg_batch.append(seg)
if len(src_batch) == args.batch_size:
yield torch.stack(src_batch, 0), \
torch.LongTensor(seg_batch)
src_batch, seg_batch = [], []
if len(src_batch) > 0:
yield torch.stack(src_batch, 0), \
torch.LongTensor(seg_batch)
def main():
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
infer_opts(parser)
parser.add_argument("--labels_num", type=int, required=True,
help="Number of prediction labels.")
tokenizer_opts(parser)
parser.add_argument("--output_logits", action="store_true", help="Write logits to output file.")
parser.add_argument("--output_prob", action="store_true", help="Write probabilities to output file.")
args = parser.parse_args()
# Load the hyperparameters from the config file.
args = load_hyperparam(args)
# Build tokenizer.
args.tokenizer = str2tokenizer["virtual"](args)
# Build classification model and load parameters.
args.soft_targets, args.soft_alpha = False, False
model = Classifier(args)
model = load_model(model, args.load_model_path)
# For simplicity, we use DataParallel wrapper to use multiple GPUs.
args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(args.device)
if torch.cuda.device_count() > 1:
print("{} GPUs are available. Let's use them.".format(torch.cuda.device_count()))
model = torch.nn.DataParallel(model)
model.eval()
with open(args.prediction_path, mode="w", encoding="utf-8") as f:
f.write("label")
if args.output_logits:
f.write("\t" + "logits")
if args.output_prob:
f.write("\t" + "prob")
f.write("\n")
for i, (src_batch, seg_batch) in enumerate(data_loader(args, args.test_path)):
src_batch = src_batch.to(args.device)
seg_batch = seg_batch.to(args.device)
with torch.no_grad():
_, logits = model(src_batch, None, seg_batch)
pred = torch.argmax(logits, dim=1)
pred = pred.cpu().numpy().tolist()
prob = nn.Softmax(dim=1)(logits)
logits = logits.cpu().numpy().tolist()
prob = prob.cpu().numpy().tolist()
for j in range(len(pred)):
f.write(str(pred[j]))
if args.output_logits:
f.write("\t" + " ".join([str(v) for v in logits[j]]))
if args.output_prob:
f.write("\t" + " ".join([str(v) for v in prob[j]]))
f.write("\n")
if __name__ == "__main__":
main()
|