def main(args):
#hyper parameter
end_epoch = 200
lr = 0.001
beta1 = 0.5
beta2 = 0.99
gpu = 0
#set model
model = MyNet()
#set GPU or CPU
if gpu >= 0 and torch.cuda.is_available():
device = 'cuda:{}'.format(gpu)
else:
device = 'cpu'
model.to(device)
model.load_state_dict(
torch.load("models/" + args.model + "/" + str(args.epoch - 1) +
".pth"))
model.eval()
criteria = nn.CrossEntropyLoss()
dataset = MyDataset("data/", is_train=False)
test_loader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True)
with torch.no_grad():
epoch_loss = 0
epoch_acc = 0
if not os.path.exists("result/" + args.model + "/true_favo/"):
os.makedirs("result/" + args.model + "/true_favo/")
os.makedirs("result/" + args.model + "/true_no_favo/")
os.makedirs("result/" + args.model + "/false_favo/")
os.makedirs("result/" + args.model + "/false_no_favo/")
for i, data in enumerate(test_loader):
print("\riteration: {}".format(i), end="")
inputs, labels = data
outputs = model(inputs.to(device))
_, preds = torch.max(outputs.data, 1)
loss = criteria(outputs, labels.to(device))
favo_preds = preds.to('cpu')
if favo_preds == labels:
if favo_preds == 1:
save_image(
inputs, "result/" + args.model + "/true_favo/" +
str(i) + ".jpg")
elif favo_preds == 0:
save_image(
inputs, "result/" + args.model + "/true_no_favo/" +
str(i) + ".jpg")
else:
if favo_preds == 1:
save_image(
inputs, "result/" + args.model + "/false_favo/" +
str(i) + ".jpg")
elif favo_preds == 0:
save_image(
inputs, "result/" + args.model + "/false_no_favo/" +
str(i) + ".jpg")
epoch_loss += loss.data.to('cpu') * inputs.size(0)
epoch_acc += torch.sum(favo_preds == labels.data)
epoch_loss /= len(test_loader)
epoch_acc = epoch_acc / float(len(test_loader))
print("[Loss: {}] [Acc: {}]".format(epoch_loss, epoch_acc))
def main(args):
#hyper parameter
end_epoch = 100
lr = 0.001
beta1 = 0.5
beta2 = 0.99
gpu = 0
#set model
model = MyNet()
#set GPU or CPU
if gpu >= 0 and torch.cuda.is_available():
device = 'cuda:{}'.format(gpu)
else:
device = 'cpu'
model.to(device)
#print params
params = 0
for p in model.parameters():
if p.requires_grad:
params += p.numel()
print(params)
print(model)
criteria = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
betas=(beta1, beta2))
dataset = MyDataset("data/", is_train=True)
train_loader = torch.utils.data.DataLoader(dataset,
batch_size=4,
shuffle=True)
for epoch in range(end_epoch):
epoch_loss = 0
epoch_acc = 0
for i, data in enumerate(train_loader):
print("\repoch: {} iteration: {}".format(epoch, i), end="")
inputs, labels = data
optimizer.zero_grad()
outputs = model(inputs.to(device))
_, preds = torch.max(outputs.data, 1)
loss = criteria(outputs, labels.to(device))
loss.backward()
optimizer.step()
epoch_loss += loss.data.to('cpu') * inputs.size(0)
epoch_acc += torch.sum(preds.to('cpu') == labels.data)
epoch_loss /= len(train_loader) * 4
epoch_acc = epoch_acc / float(len(train_loader) * 4)
print("[epoch: {}] [Loss: {:.4f}] [Acc: {:.4f}]".format(
epoch, epoch_loss, epoch_acc))
if (epoch + 1) % 10 == 0:
if not os.path.exists("models/" + args.model):
os.makedirs("models/" + args.model)
torch.save(model.state_dict(),
"models/" + args.model + "/" + str(epoch) + ".pth")
def main(self):
"""
训练接口主函数，完成整个训练流程
1. 创建训练集和验证集的DataLoader类
2. 初始化带训练的网络
3. 选择合适的优化器
4. 训练并验证指定个epoch，保存其中评价指标最好的模型，并打印训练过程信息
5. TODO: 可视化训练过程信息
"""
opts = self.opts
if not os.path.exists(opts.checkpoints_dir):
os.mkdir(opts.checkpoints_dir)
random_seed = opts.random_seed
train_dataset = MyDataset(opts.dataset_dir,
seed=random_seed,
mode="train",
train_val_ratio=0.9)
val_dataset = MyDataset(opts.dataset_dir,
seed=random_seed,
mode="val",
train_val_ratio=0.9)
train_loader = DataLoader(train_dataset,
opts.batch_size,
shuffle=True,
num_workers=opts.num_workers)
val_loader = DataLoader(val_dataset,
batch_size=1,
shuffle=False,
num_workers=opts.num_workers)
num_train = len(train_dataset)
num_val = len(val_dataset)
if opts.pretrain is None:
model = MyNet()
else:
model = torch.load(opts.pretrain)
if opts.use_GPU:
model.to(opts.GPU_id)
optimizer = torch.optim.SGD(model.parameters(),
lr=opts.lr,
momentum=0.9,
weight_decay=opts.weight_decay)
# optimizer = torch.optim.Adam(model.parameters(), lr=opts.lr, weight_decay=opts.weight_decay)
best_metric = 1000000
for e in range(opts.start_epoch, opts.epoch + 1):
t = time.time()
self.__train(model, train_loader, optimizer, e, num_train, opts)
t2 = time.time()
print("Training consumes %.2f second\n" % (t2 - t))
with open(os.path.join(opts.checkpoints_dir, "log.txt"),
"a+") as log_file:
log_file.write("Training consumes %.2f second\n" % (t2 - t))
if e % opts.save_freq == 0 or e == opts.epoch + 1:
# t = time.time()
# metric = self.__validate(model, val_loader, e, num_val, opts)
# t2 = time.time()
# print("Validation consumes %.2f second\n" % (t2 - t))
# with open(os.path.join(opts.checkpoints_dir, "log.txt"), "a+") as log_file:
# log_file.write("Validation consumes %.2f second\n" % (t2 - t))
# if best_metric>metric:
# best_metric = metric
# print("Epoch %d is now the best epoch with metric %.4f\n"%(e, best_metric))
# with open(os.path.join(opts.checkpoints_dir, "log.txt"), "a+") as log_file:
# log_file.write("Epoch %d is now the best epoch with metric %.4f\n"%(e, best_metric))
self.__save_model(model, e, opts)
