print(test_2_list)
data1 = myData(datapath, train_list)
data2 = myData(datapath, test_2_list)
train_data = DataLoader(dataset=data1, batch_size=48, shuffle=True)
test_data = DataLoader(dataset=data2, batch_size=1, shuffle=True)
data = {'train': train_data, 'test': test_data}
data_len = {'train': len(train_list), 'test': len(test_2_list)}
cnn = CNN()
if torch.cuda.is_available():
cnn = cnn.cuda()
print('正在使用GPU...')
else:
print('正在使用CPU...')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
cnn.apply(weights_init)
optimizer = torch.optim.Adam(cnn.parameters(), lr=lr)
loss_func = nn.CrossEntropyLoss()
train_model(cnn, num_epochs, data, device, optimizer, loss_func, data_len,
savepath)
def train_model(load_model=None, num_features=16):
experiment = {
'seed': 1234,
'model_kernel': 3,
'model_num_features': 10,
'model_dropout_rate': 0.01,
'train_class_weight': [.18, .02, .80], # BALL, ROBOT, OTHER
'train_learn_rate': 1e-2, # 1e-3,
'train_weight_decay': 1e-9,
'train_batch_size': 8, # 32, # 80, # 20,
'train_epochs': 7,
'colour_jitter': [0.0, 0.0, 0.0, 0.0], # brightness, contrast, saturation, hue
'output_folder': 'nam-training3',
}
# Save directory
output_folder = experiment['output_folder']
if not os.path.exists(output_folder):
os.makedirs(output_folder)
# test multiple times for statistical significance
valid_final_losses = []
test_final_losses = []
test_precision = []
test_recall = []
for i in range(1):
if load_model is not None:
print(f'Loading previously trained model: {load_model}')
model = CNN(kernel=3, num_features=num_features)
model.load_state_dict(torch.load(load_model))
else:
model = CNN(
kernel=experiment['model_kernel'],
num_features=experiment['model_num_features'],
dropout=experiment['model_dropout_rate'])
model.apply(init_weights)
trainer = Trainer(model,
learn_rate=experiment['train_learn_rate'],
weight_decay=experiment['train_weight_decay'],
batch_size=experiment['train_batch_size'],
epochs=experiment['train_epochs'],
colour_jitter=experiment['colour_jitter'],
output_folder=experiment['output_folder'],
seed=experiment['seed'],
class_weights=experiment['train_class_weight'])
trainer.train()
valid_final_losses.append(trainer.valid_losses[-2])
test_final_losses.append(trainer.valid_losses[-1])
test_precision.append(trainer.precision)
test_recall.append(trainer.recall)
print(valid_final_losses)
print(test_final_losses)
print('valid mean loss:', np.mean(valid_final_losses), ', std:', np.std(valid_final_losses))
print('test mean loss: ', np.mean(test_final_losses), ', std:', np.std(test_final_losses))
print(test_precision)
print(test_recall)
print('test precision:', np.mean(test_precision), ', std:', np.std(test_precision))
print('test recall: ', np.mean(test_recall), ', std:', np.std(test_recall))
print(experiment)
torch.save(model.state_dict(), 'outputs/model')
