【pytorch使用tfrecord数据】

pytorch读取tfrecord文件

import os
import sys

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision
from torchvision import transforms, datasets
from tqdm import tqdm
from tfrecord.torch.dataset import TFRecordDataset

from model import resnet18


def parse_tfrecord(single_record):  #解码tfrecord文件并将其转化成训练时的张量，这个其实是对数据transform
    image, label = single_record    #和tensorflow解码类似
    image = torch.tensor(single_record["image/encoded"]) #先将其转化为向量
    
    label = torch.tensor(single_record["image/label"]).squeeze() 
    #label读出[[label1], [label2],...]，如果不降维，你每次取label就直接是一个元组[label1]，无法进行训练
    #降维之后就是[label1, label2,...]
    
    image = torchvision.io.decode_jpeg(image).float()
    #将image读出来重组成jpeg，操作和tensorflow类似，.float()是因为网络权重是float类型，两者必须相同
    
    return (image, label)
    

def main():
    device = torch.device("cuda:4" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))
    batch_s