def read_data(batchsz=32):
# 加载训练数据
train_data = mydataset.MyDataset(
'./data/train.txt',
transform=transforms.Compose([
# 数据增强：随机水平翻转
transforms.RandomHorizontalFlip(),
# 数据增强：随机垂直翻转
transforms.RandomVerticalFlip(),
# 转换成张量
transforms.ToTensor(),
# 数据增强，像素归一化，均值 / 方差
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]))
train_loader = DataLoader(train_data, batch_size=batchsz, shuffle=True)
# 加载测试数据
test_data = mydataset.MyDataset(
'./data/test.txt',
transform=transforms.Compose([
# 数据增强：随机水平翻转
transforms.RandomHorizontalFlip(),
# 数据增强：随机垂直翻转
transforms.RandomVerticalFlip(),
# 转换成张量
transforms.ToTensor(),
# 数据增强，像素归一化，均值 / 方差
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]))
test_loader = DataLoader(test_data, batch_size=batchsz, shuffle=True)
x, label = iter(train_loader).next()
print('x:', x.shape, 'label:', label.shape)
print("数据加载完毕...")
return train_data, train_loader, test_data, test_loader
def load_data():
"""Load the data: The detail folders contain a '1 bee' folder with 78x78 grayscale PNG images
of bees in honeycombs and a '0 notbee' folder with negative examples."""
b_size = 200
sigma = 15
transform_train = transforms.Compose([
transforms.RandomOrder([
transforms.RandomRotation(180),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ColorJitter(brightness=0.5, contrast=0.5)
]),
transforms.CenterCrop(54),
transforms.Resize(32),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor()
])
transform_test = transforms.Compose([
transforms.CenterCrop(54),
transforms.Resize(32),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor()
])
trainset = beeimagefolder.BeeImageFolder(
root='../Videos/detail_training', valid_classes=['0 notbee', '1 bee'])
translated_trainset = mydataset.MyDataset(dataset=trainset,
sigma=sigma,
transform=transform_train)
testset = beeimagefolder.BeeImageFolder(
root='../Videos/detail_test',
valid_classes=['0 notbee', '1 bee'],
transform=transform_test)
print('class_to_idx trainset: {}\t class_to_idx testset: {}'.format(
trainset.class_to_idx, testset.class_to_idx))
trainloader = torch.utils.data.DataLoader(translated_trainset,
num_workers=2,
batch_size=b_size,
shuffle=True)
testloader = torch.utils.data.DataLoader(testset)
classes = ('notbee', 'bee')
return trainloader, testloader, classes, b_size, sigma
def read_data(batchsz=32):
# 加载训练数据
test_data = mydataset.MyDataset(
'./data/test.txt',
transform=transforms.Compose([
# 转换成张量
transforms.ToTensor(),
# 数据增强，像素归一化，均值 / 方差
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]))
test_loader = DataLoader(test_data, batch_size=batchsz, shuffle=False)
x, label, fn = iter(test_loader).next()
print('x:', x.shape, 'label:', label.shape, 'fn:', len(fn))
print("数据加载完毕...")
return test_data, test_loader
def __init__(self): # 初始化
# 定义训练运行在什么设备上：cpu或者gpu
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
# 创建数据集Dataset，然后把Dataset扔进torch的dataloader里
self.dataset = mydataset.MyDataset()
self.dataloader = DataLoader(
self.dataset,
batch_size=2,
shuffle=True,
num_workers=1,
pin_memory=True,
)
# 构造网络
self.net = model.MaskRCNN()
# 训练超参数的设置
self.optimizer = torch.optim.AdamW() # 优化器，图省事用的Adam系，因为迷信所以用的AdamW
self.criterion = None
os.mkdir('log')
if config.use_log and os.path.exists(log_filename):
os.remove(log_filename)
if config.experiment is None:
config.experiment = 'expr'
if not os.path.exists(config.experiment):
os.mkdir(config.experiment)
config.manualSeed = random.randint(1, 10000) # fix seed
print("Random Seed: ", config.manualSeed)
random.seed(config.manualSeed)
np.random.seed(config.manualSeed)
torch.manual_seed(config.manualSeed)
# cudnn.benchmark = True
train_dataset = mydataset.MyDataset(info_filename=config.train_infofile)
assert train_dataset
if not config.random_sample:
sampler = mydataset.randomSequentialSampler(train_dataset,
config.batchSize)
else:
sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=config.batchSize,
shuffle=True,
sampler=sampler,
num_workers=int(config.workers),
collate_fn=mydataset.alignCollate(
imgH=config.imgH,
imgW=config.imgW,
keep_ratio=config.keep_ratio))
print "\nParameters:"
for attr, value in sorted(args.__dict__.items()):
print "\t%s=%s" % (attr.upper(), value)
# model
cnn = model.Model(args)
# load initial model
if args.snapshot is not None:
print "Loading model from %s..." % (args.snapshot)
cnn.load_state_dict(torch.load(args.snapshot))
if (args.train_predict) or (args.train):
# data
print "loading dataset..."
dataset = mydataset.MyDataset(args)
validset = mydataset.MyDataset(args, valid_set=True)
# train
print "starting training..."
train.train(dataset, validset, cnn, args)
if (args.train_predict) or (args.predict):
print "loading testset..."
testset = mydataset.MyTestset(args)
print "predicting..."
outputs = train.predict(testset, cnn, args)
#
fout = file(args.pred_outfile, "w")
for i in outputs:
alphabet = []
with open(config.character_file, 'r', encoding='utf-8') as file:
alphabet = [s.strip('\n') for s in file.readlines()]
config.alphabet = [ord(ch) for ch in alphabet]
config.nclass = len(config.alphabet) + 1
print('nlacss', config.nclass)
config.manualSeed = random.randint(1, 10000) # fix seed
print("Random Seed: ", config.manualSeed)
random.seed(config.manualSeed)
np.random.seed(config.manualSeed)
torch.manual_seed(config.manualSeed)
# cudnn.benchmark = True
train_dataset = mydataset.MyDataset(info_filename=config.train_infofile,
character_file=config.character_file,
image_fold=config.train_data_fold)
assert train_dataset
if not config.random_sample:
sampler = mydataset.randomSequentialSampler(train_dataset,
config.batchSize)
else:
sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.batchSize,
shuffle=True,
sampler=sampler,
# num_workers=int(config.workers),
collate_fn=mydataset.alignCollate(imgH=config.imgH,
imgW=config.imgW,