# resnet18を取得
Net = OPN(pretrained=args.use_pretrained_model)
criterion = torch.nn.CrossEntropyLoss() # Loss関数を定義
optimizer = torch.optim.Adam(Net.parameters(),
lr=args.learning_rate) # 重み更新方法を定義
current_epoch = 0
# ログファイルの生成
if not args.no_reset_log_file:
with open(log_train_path, mode='w') as f:
f.write('epoch,loss,accuracy,time,learning_rate\n')
# CUDA環境の有無で処理を変更
if args.use_cuda:
criterion = criterion.cuda()
Net = torch.nn.DataParallel(Net.cuda())
device = 'cuda'
else:
device = 'cpu'
# モデルの読み込み
if args.model_load_path:
checkpoint = torch.load(args.model_load_path)
Net.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
if args.load_epoch_num:
current_epoch = checkpoint['epoch']
print('complete load model')
# 訓練を行う
criterion = nn.CrossEntropyLoss() # Loss関数を定義
optimizer = torch.optim.Adam(Net.parameters(),
lr=args.learning_rate) # 重み更新方法を定義
current_epoch = 0
# ログファイルの生成
if not args.no_reset_log_file:
with open(log_train_path, mode='w') as f:
f.write('epoch,loss,accuracy,time,learning_rate\n')
with open(log_test_path, mode='w') as f:
f.write('epoch,loss,accuracy,time,learning_rate\n')
# CUDA環境の有無で処理を変更
if args.use_cuda:
criterion = criterion.cuda()
Net = nn.DataParallel(Net.cuda())
device = 'cuda'
else:
device = 'cpu'
if args.model_load_path:
checkpoint = torch.load(args.model_load_path)
Net.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
if args.load_epoch_num:
current_epoch = checkpoint['epoch']
print('complete load model')
# 訓練を行う
def train(inputs, labels):