# 检查cuda是否可用
# 生成log
now_time = datetime.now()
time_str = datetime.strftime(now_time, '%m-%d_%H-%M-%S')
log_path = os.path.join(os.getcwd(), "log")
log_dir = os.path.join(log_path, time_str)
if not os.path.exists(log_dir):
os.makedirs(log_dir)
writer = SummaryWriter(log_dir)
# ---------------------搭建网络--------------------------
cnn = Net() # 创建CNN, 输入全连接层维数
print("Total number of paramerters in networks is {} ".format(sum(x.numel() for x in cnn.parameters())))
cnn = cnn.double()
# --------------------设置损失函数和优化器----------------------
optimizer = optim.Adam(cnn.parameters(), lr=0.001) # lr:(de fault: 1e-3)优化器
criterion = nn.CrossEntropyLoss() # 损失函数
scheduler = torch.optim.lr_scheduler.StepLR(
optimizer, step_size=EPOCH / 2, gamma=0.5) # 设置学习率下降策略
# --------------------训练------------------------------
# 使用GPU
cnn = cnn.cuda()
print('=========STAGE1==============')
for epoch in range(1):
loss_sigma = 0.0 # 记录一个epoch的loss之和
correct = 0.0
total = 0.0
