def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch Vec2Color Example')
parser.add_argument('--batch-size', type=int, default=64, metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=1500, metavar='N',
help='number of epochs to train (default: 1500)')
parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
help='learning rate (default: 1.0)')
parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=False,
help='For Saving the current Model')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
file_names = ('capitalize', 'lower', 'upper', 'title')
x_df = pd.concat([pd.read_csv('doc2color/data/{}.csv'.format(file_name)) for file_name in file_names])
y_df = pd.concat([pd.read_csv('doc2color/data/rgb.csv')] * len(file_names))
tensor_x = torch.stack([torch.from_numpy(np.array(i)) for i in x_df.values.astype(np.float32)])
tensor_y = torch.stack([torch.from_numpy(np.array(i)) for i in y_df.values.astype(np.float32) / 255.0])
x_train, x_test, y_train, y_test = train_test_split(
tensor_x, tensor_y, test_size=0.01, random_state=args.seed)
train_dataset = torch.utils.data.TensorDataset(x_train, y_train)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size, shuffle=True, **kwargs)
test_dataset = torch.utils.data.TensorDataset(x_test, y_test)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.test_batch_size, shuffle=True, **kwargs)
model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
test(args, model, device, test_loader)
scheduler.step()
if args.save_model:
torch.save(model.state_dict(), "doc2color/pt_objects/vec2color.pt")
def main(args):
train_loader, val_loader = create_dataloaders(args.batch_size)
model = Net().to(device)
optim = torch.optim.Adam(model.parameters())
lr_schedule = torch.optim.lr_scheduler.ReduceLROnPlateau(optim,
patience=1,
verbose=True)
criterion = torch.nn.CrossEntropyLoss()
best_accuracy = 0
for epoch in range(1, args.epochs + 1):
model.train()
train_loss, train_accuracy = do_epoch(model,
train_loader,
criterion,
optim=optim)
model.eval()
with torch.no_grad():
val_loss, val_accuracy = do_epoch(model,
val_loader,
criterion,
optim=None)
tqdm.write(
f'EPOCH {epoch:03d}: train_loss={train_loss:.4f}, train_accuracy={train_accuracy:.4f} '
f'val_loss={val_loss:.4f}, val_accuracy={val_accuracy:.4f}')
if val_accuracy > best_accuracy:
print('Saving model...')
best_accuracy = val_accuracy
torch.save(model.state_dict(), 'trained_models/source.pt')
lr_schedule.step(val_loss)
def run():
print("CUDA is available: {}".format(torch.cuda.is_available()))
data_transform = transforms.Compose(
[Rescale(250), CenterCrop(224),
Normalize(), ToTensor()])
# loader will split datatests into batches witht size defined by batch_size
train_loader = initialize_train_loader(data_transform)
test_loader = initialize_test_loader(data_transform)
model_id = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime())
# instantiate the neural network
net = Net()
net.to(device=device)
summary(net, (1, 224, 224))
# define the loss function using SmoothL1Loss
criterion = nn.SmoothL1Loss()
# define the params updating function using Adam
optimizer = optim.Adam(net.parameters(), lr=0.001)
loss_logger = []
for i in range(1, epochs + 1):
model_name = 'model-{}-epoch-{}.pt'.format(model_id, i)
# train all data for one epoch
train(net, criterion, optimizer, i, train_loader, model_id,
loss_logger)
# evaludate the accuracy after each epoch
evaluate(net, criterion, i, test_loader)
# save model after every 5 epochs
# https://discuss.pytorch.org/t/loading-a-saved-model-for-continue-training/17244/3
# https://github.com/pytorch/pytorch/issues/2830
# https://pytorch.org/tutorials/beginner/saving_loading_models.html
if i % 5 == 1:
torch.save(
{
'epoch': i,
'model': net.state_dict(),
'optimizer': optimizer.state_dict(),
'loss_logger': loss_logger,
}, model_dir + model_name)
print("Finished training!")
model_name = 'model-{}-final.pt'.format(model_id)
torch.save(
{
'epoch': epochs,
'model': net.state_dict(),
'optimizer': optimizer.state_dict(),
'loss_logger': loss_logger,
}, model_dir + model_name)
def initialize_networks(self):
self.networks = []
self.optimizers = []
for i in range(self.n_models):
model = Net()
model.apply(self.init_weights)
self.networks.append(model)
self.optimizers.append(
optim.SGD(model.parameters(),
lr=self.learning_rate,
momentum=self.momentum))
self.criterion = nn.CrossEntropyLoss()
def load_model():
"""
Get the newest checkpoint (best one) in the checkpoint directory
For more detail about torch's load/save progress, see:
https://stackoverflow.com/questions/42703500/best-way-to-save-a-trained-model-in-pytorch
:param filename:
:return: model (with training progress, can be used for either inference or further training)
"""
model = Net(int(args['in_size']), int(args['hidden_size']),
int(args['out_size']), int(args['nb_lstm_layers']),
int(args['batch_size']), int(bool(args['bidirectional'])))
newest_dir = [
xxx for xxx in subprocess.check_output("ls -t checkpoints/".split(
" ")).decode("utf8").strip().split("\n")
if os.path.isdir("checkpoints/" + xxx)
][0] # | head -n 1"])
print(newest_dir)
newest_pth = [
"checkpoints/{}/{}".format(newest_dir, xxx)
for xxx in subprocess.check_output("ls -t checkpoints/{}/".format(
newest_dir).split(" ")).decode("utf8").strip().split("\n")
if os.path.isfile("checkpoints/{}/{}".format(newest_dir, xxx))
]
list_file = {
get_epoch_from_filename(filename)[0]:
get_epoch_from_filename(filename)[1]
for filename in newest_pth
}
best_epoch = max(list_file.keys())
best_loss = list_file[best_epoch]
print(best_epoch, best_loss)
state = torch.load("checkpoints/{}/checkpoint_epoch{}_{}".format(
newest_dir, best_epoch, best_loss),
map_location='cpu')
model.load_state_dict(state['state_dict'])
optimizer = optim.RMSprop(model.parameters())
optimizer.load_state_dict(state['optimizer'])
return model, optimizer
def main():
use_cuda = torch.cuda.is_available() and args.cuda
device = torch.device('cuda' if use_cuda else 'cpu')
print('Loading model ...')
model = Net(device)
if args.load is not None:
print('Loading checkpoint ...')
model.load_state_dict(torch.load(args.load))
if use_cuda:
model.cuda()
print('Loading data ...')
data_loader = get_data_loader(args.data_root, args.stage, args.batch_size)
print('Preparation done')
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.999))
criterion = nn.CrossEntropyLoss()
if args.stage == 'train':
train(model, data_loader, optimizer, criterion)
def main(args):
epochs = args.epochs
batch_size = args.batch_size
val_frac = args.val_frac
lr = args.lr
momentum = args.momentum
total_it = args.total_it
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
dataset_train = HomographyDataset(val_frac=val_frac, mode='train')
dataset_eval = HomographyDataset(val_frac=val_frac, mode='eval')
dataloader_train = DataLoader(dataset_train,
batch_size=batch_size,
shuffle=True,
num_workers=4)
dataloader_eval = DataLoader(dataset_eval,
batch_size=batch_size,
shuffle=False,
num_workers=2)
net = Net()
net.to(device)
criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), momentum=momentum, lr=lr)
log_every = len(dataloader_train) // 2
n_it, lr_it = 0, 0
while n_it < total_it:
train(dataloader_train, device, net, criterion, optimizer)
n_it += len(dataloader_train)
test(dataloader_eval, device, net, criterion, n_it)
if lr_it >= 30000:
d = optimizer.state_dict()
d['param_groups'][0]['lr'] /= 10.0
optimizer.load_state_dict(d)
lr_it = 0
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--datapath",
type=str,
default="data/3droad.mat",
help="dataset file")
parser.add_argument('--epochs',
type=int,
default=30,
help='number of epochs to train')
parser.add_argument('--batch_size',
type=int,
default=1024,
help='input batch size for training')
parser.add_argument('--val_batch_size',
type=int,
default=1024,
help='input batch size for valing')
# Training Settings
parser.add_argument('--arch_file',
type=str,
default=None,
help='name of file which defines the architecture')
parser.add_argument('--arch_name',
type=str,
default='LeNet5',
help='name of the architecture')
parser.add_argument('--arch_args',
type=json.loads,
default=None,
help='[JSON] arguments for the architecture')
parser.add_argument('--optim_name',
type=str,
default=VIOptimizer.__name__,
help='name of the optimizer')
parser.add_argument('--optim_args',
type=json.loads,
default=None,
help='[JSON] arguments for the optimizer')
parser.add_argument('--curv_args',
type=json.loads,
default=dict(),
help='[JSON] arguments for the curvature')
parser.add_argument('--fisher_args',
type=json.loads,
default=dict(),
help='[JSON] arguments for the fisher')
parser.add_argument('--scheduler_name',
type=str,
default=None,
help='name of the learning rate scheduler')
parser.add_argument('--scheduler_args',
type=json.loads,
default=None,
help='[JSON] arguments for the scheduler')
# Options
parser.add_argument('--no_cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, help='random seed')
parser.add_argument(
'--log_interval',
type=int,
default=20,
help='how many batches to wait before logging training status')
parser.add_argument(
'--checkpoint_interval',
type=int,
default=50,
help='how many epochs to wait before logging training status')
parser.add_argument('--out',
type=str,
default='results/regression/',
help='dir to save output files')
parser.add_argument('--config', default=None, help='config file path')
parser.add_argument("--log_name",
default=None,
required=True,
type=str,
help="log name")
args = parser.parse_args()
dict_args = vars(args)
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
if args.config is not None:
with open(args.config) as f:
config = json.load(f)
dict_args.update(config)
if not os.path.isfile(args.datapath):
print('Downloading \'3droad\' UCI dataset...')
urllib.request.urlretrieve(
'https://www.dropbox.com/s/f6ow1i59oqx05pl/3droad.mat?dl=1',
args.datapath)
data = loadmat(args.datapath)['data']
X = data[:, :-1]
X = X - X.min(0)[0]
X = 2 * (X / X.max(0)[0]) - 1
y = data[:, -1]
train_val_n = int(floor(0.8 * len(X)))
X_train_val = X[:train_val_n]
y_train_val = y[:train_val_n]
X_test = X[train_val_n:, :]
y_test = y[train_val_n:]
X_train, X_val, y_train, y_val = train_test_split(X_train_val,
y_train_val,
test_size=0.2)
dtype = torch.float32
X_train = torch.tensor(X_train, dtype=dtype, requires_grad=False)
y_train = torch.tensor(y_train, dtype=dtype, requires_grad=False)
X_val = torch.tensor(X_val, dtype=dtype, requires_grad=False)
y_val = torch.tensor(y_val, dtype=dtype, requires_grad=False)
X_test = torch.tensor(X[train_val_n:, :], dtype=dtype, requires_grad=False)
y_test = torch.tensor(y[train_val_n:], dtype=dtype, requires_grad=False)
train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True)
val_dataset = TensorDataset(X_val, y_val)
val_loader = DataLoader(val_dataset,
batch_size=args.val_batch_size,
shuffle=True)
test_dataset = TensorDataset(X_test, y_test)
test_loader = DataLoader(test_dataset,
batch_size=args.val_batch_size,
shuffle=True)
data_dim = X_train.size(-1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net(data_dim).to(device).float()
optim_kwargs = {} if args.optim_args is None else args.optim_args
if args.optim_name == SecondOrderOptimizer.__name__:
optimizer = SecondOrderOptimizer(model,
**config["optim_args"],
curv_kwargs=config["curv_args"])
elif args.optim_name == VIOptimizer.__name__:
optimizer = VIOptimizer(model,
dataset_size=len(train_loader.dataset),
seed=args.seed,
**config["optim_args"],
curv_kwargs=config["curv_args"])
else:
modules = import_module("optimizers")
optim_class = getattr(modules, args.optim_name)
optimizer = optim_class(model.parameters())
if args.scheduler_name is None:
scheduler = None
else:
scheduler_class = getattr(torchsso.optim.lr_scheduler,
args.scheduler_name, None)
if scheduler_class is None:
scheduler_class = getattr(torch.optim.lr_scheduler,
args.scheduler_name)
scheduler_kwargs = config["scheduler_args"]
scheduler = scheduler_class(optimizer, **scheduler_kwargs)
log_file_name = "log_" + args.log_name
logger = Logger(args.out, log_file_name)
logger.start()
# train
epochs = 120
model.train()
print("=========== Start ===========")
for i in range(args.epochs):
losses = 0
for minibatch_i, (x_batch, y_batch) in enumerate(train_loader):
x_batch = x_batch.to(device)
y_batch = y_batch.to(device)
"""
def closure():
optimizer.zero_grad()
output = model(x_batch)
loss = F.mse_loss(output, y_batch, reduction="sum").float()
loss.backward()
return loss, output
"""
optimizer.zero_grad()
output = model(x_batch)
loss = F.mse_loss(output, y_batch, reduction="sum").float()
loss.backward()
optimizer.step()
losses += loss.item()
if (minibatch_i + 1) % args.log_interval == 0:
print(
"Train Epoch: {} batch idx: {} elapsed time: {:.1f}s MSE: {}"
.format(i + 1, minibatch_i + 1, logger.elapsed_time,
loss.item() / args.batch_size))
losses = losses / len(train_loader.dataset)
val_mse = validate(model, device, val_loader, optimizer, "Eval")
iteration = (i + 1) * len(train_loader)
log = {
"epoch": i + 1,
"iteration": iteration,
"mse": losses,
"val_mse": val_mse,
"lr": optimizer.param_groups[0]["lr"],
"momentum": optimizer.param_groups[0].get("momentum", 0)
}
logger.write(log)
if i % args.checkpoint_interval == 0 or i + 1 == args.epochs:
path = os.path.join(args.out,
"epoch{}_{}.ckpt".format(i + 1, args.log_name))
data = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"epoch": i + 1
}
torch.save(data, path)
print("=========== Test ===========")
test_mse = validate(model, device, test_loader, optimizer, "Test")
log = {"test_mse": test_mse}
logger.write(log)
def main():
# Training settings
parser = argparse.ArgumentParser(description='Grid Lenet Trainer')
parser.add_argument('--batch-size',
type=int,
default=64,
metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--epochs',
type=int,
default=10,
metavar='N',
help='number of iterations to train (default: 100)')
parser.add_argument('--epoch_len', type=int, default=10)
parser.add_argument('--lr',
type=float,
default=0.001,
metavar='LR',
help='learning rate (default: 0.01)')
parser.add_argument('--no-cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument('--deterministic',
type=str2bool,
help='make dataloader deterministic')
parser.add_argument(
'--log-interval',
type=int,
default=10,
metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument(
'--num-samples',
type=int,
default=5,
help='the number of samples for renders predicitons on')
parser.add_argument('--dataset', default='./data', help='path to dataset')
parser.add_argument('--restore',
type=str2bool,
help='restore from checkpoint if available')
parser.add_argument('--restore-from',
help='checkpoint name to use instead of default one')
parser.add_argument('--exp-dir',
default='experiments',
help='place for storing experiment-related data')
parser.add_argument('--name',
required=True,
help='Network name, for saving checkpoints and logs')
parser.add_argument('--wd', default=0, type=float, help='weight-decay')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if use_cuda:
torch.cuda.manual_seed(args.seed)
grid_dataset = GridDataset(path=args.dataset)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
train_loader = DataLoader(dataset=grid_dataset,
batch_size=args.batch_size,
shuffle=not args.deterministic,
**kwargs)
model = Net().to(device)
model.device = device
opt = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
if args.restore:
epoch = load_checkpoint(args, model, opt)
else:
epoch = 1
writer = SummaryWriter(get_log_dir(args))
samples = []
stop_sample_collection = False
while not stop_sample_collection:
for i in range(len(grid_dataset)):
samples.append(grid_dataset[i])
if len(samples) >= args.num_samples:
stop_sample_collection = True
break
while epoch <= args.epochs:
train(args, model, train_loader, opt, epoch, writer=writer)
if args.log_interval:
render_predictions(model, samples, epoch,
get_sample_path(args, 'train'))
save_checkpoint(args, model, opt, epoch)
epoch += 1
# def show_all_keypoints(image, predicted_key_pts, gt_pts=None):
# """Show image with predicted keypoints"""
# # image is grayscale
# plt.imshow(image, cmap='gray')
# plt.scatter(predicted_key_pts[:, 0], predicted_key_pts[:, 1], s=20, marker='.', c='m')
# # plot ground truth points as green pts
# if gt_pts is not None:
# plt.scatter(gt_pts[:, 0], gt_pts[:, 1], s=20, marker='.', c='g')
## TODO: Define the loss and optimization
import torch.optim as optim
criterion = nn.MSELoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
def train_net(n_epochs):
# prepare the net for training
net.train()
for epoch in range(n_epochs): # loop over the dataset multiple times
running_loss = 0.0
# train on batches of data, assumes you already have train_loader
for batch_i, data in enumerate(train_loader):
# get the input images and their corresponding labels
images = data['image'].to(device)
def train(training_opt: TrainingOptions, device_opt: DeviceOptions) -> Net:
if device_opt.cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
torch.manual_seed(training_opt.seed)
device = torch.device("cuda" if device_opt.cuda else "cpu")
print('===> Loading datasets')
train_set = get_train_set(training_opt.upscale_factor)
test_set = get_test_set(training_opt.upscale_factor)
training_data_loader = DataLoader(dataset=train_set,
num_workers=device_opt.threads,
batch_size=device_opt.batchSize,
shuffle=True)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=device_opt.threads,
batch_size=device_opt.testBatchSize,
shuffle=False)
print('===> Building model')
model = Net(upscale_factor=training_opt.upscale_factor).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=training_opt.lr)
training_status = st.empty()
test_status = st.empty()
def train(epoch: int):
training_status.code('Training... (epoch=%s)' % epoch)
epoch_loss = 0
for iteration, batch in enumerate(training_data_loader, 1):
input, target = batch[0].to(device), batch[1].to(device)
optimizer.zero_grad()
loss = criterion(model(input), target)
epoch_loss += loss.item()
loss.backward()
optimizer.step()
training_status.code(
"===> Epoch[{}/{}]({}/{}): Loss: {:.4f}".format(
epoch, training_opt.nEpochs, iteration,
len(training_data_loader), loss.item()))
training_status.code(
"===> Epoch {} Complete: Avg. Loss: {:.4f}".format(
epoch, epoch_loss / len(training_data_loader)))
def test():
avg_psnr = 0
with torch.no_grad():
for batch in testing_data_loader:
input, target = batch[0].to(device), batch[1].to(device)
prediction = model(input)
mse = criterion(prediction, target)
psnr = 10 * log10(1 / mse.item())
avg_psnr += psnr
test_status.code("===> Avg. PSNR: {:.4f} dB".format(
avg_psnr / len(testing_data_loader)))
for epoch in range(1, training_opt.nEpochs + 1):
train(epoch)
# test()
return model
def main(args):
source_model = Net().to(device)
source_model.load_state_dict(torch.load(args.MODEL_FILE))
source_model.eval()
set_requires_grad(source_model, requires_grad=False)
clf = source_model
source_model = source_model.feature_extractor
target_model = Net().to(device)
target_model.load_state_dict(torch.load(args.MODEL_FILE))
target_model = target_model.feature_extractor
classifier = clf.classifier
discriminator = nn.Sequential(nn.Linear(EXTRACTED_FEATURE_DIM, 64),
nn.ReLU(), nn.BatchNorm1d(64),
nn.Linear(64, 1), nn.Sigmoid()).to(device)
#half_batch = args.batch_size // 2
batch_size = args.batch_size
# X_source, y_source = preprocess_train()
X_source, y_source = preprocess_train_single(1)
source_dataset = torch.utils.data.TensorDataset(X_source, y_source)
source_loader = DataLoader(source_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1,
pin_memory=True)
X_target, y_target = preprocess_test(args.person)
target_dataset = torch.utils.data.TensorDataset(X_target, y_target)
target_loader = DataLoader(target_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=1,
pin_memory=True)
discriminator_optim = torch.optim.Adam(discriminator.parameters())
target_optim = torch.optim.Adam(target_model.parameters(), lr=3e-6)
criterion = nn.BCEWithLogitsLoss()
criterion_class = nn.CrossEntropyLoss()
best_tar_acc = test(args, clf)
final_accs = []
for epoch in range(1, args.epochs + 1):
source_loader = DataLoader(source_loader.dataset,
batch_size=batch_size,
shuffle=True)
target_loader = DataLoader(target_loader.dataset,
batch_size=batch_size,
shuffle=True)
batch_iterator = zip(loop_iterable(source_loader),
loop_iterable(target_loader))
total_loss = 0
adv_loss = 0
total_accuracy = 0
second_acc = 0
total_class_loss = 0
for _ in trange(args.iterations, leave=False):
# Train discriminator
set_requires_grad(target_model, requires_grad=False)
set_requires_grad(discriminator, requires_grad=True)
discriminator.train()
for _ in range(args.k_disc):
(source_x, source_y), (target_x, _) = next(batch_iterator)
source_y = source_y.to(device).view(-1)
source_x, target_x = source_x.to(device), target_x.to(device)
source_features = source_model(source_x).view(
source_x.shape[0], -1)
target_features = target_model(target_x).view(
target_x.shape[0], -1)
discriminator_x = torch.cat([source_features, target_features])
discriminator_y = torch.cat([
torch.ones(source_x.shape[0], device=device),
torch.zeros(target_x.shape[0], device=device)
])
preds = discriminator(discriminator_x).squeeze()
loss = criterion(preds, discriminator_y)
discriminator_optim.zero_grad()
loss.backward()
discriminator_optim.step()
total_loss += loss.item()
total_accuracy += ((preds >= 0.5).long() == discriminator_y.
long()).float().mean().item()
# Train feature extractor
set_requires_grad(target_model, requires_grad=True)
set_requires_grad(discriminator, requires_grad=False)
target_model.train()
for _ in range(args.k_clf):
_, (target_x, _) = next(batch_iterator)
target_x = target_x.to(device)
target_features = target_model(target_x).view(
target_x.shape[0], -1)
source_features = target_model(source_x).view(
source_x.shape[0], -1)
source_pred = classifier(source_features) # (batch_size, 4)
# flipped labels
discriminator_y = torch.ones(target_x.shape[0], device=device)
preds = discriminator(target_features).squeeze()
second_acc += ((preds >= 0.5).long() == discriminator_y.long()
).float().mean().item()
loss_adv = criterion(preds, discriminator_y)
adv_loss += loss_adv.item()
loss_class = criterion_class(source_pred, source_y)
total_class_loss += loss_class.item()
loss = loss_adv #+ 0.001*loss_class
target_optim.zero_grad()
loss.backward()
target_optim.step()
mean_loss = total_loss / (args.iterations * args.k_disc)
mean_adv_loss = adv_loss / (args.iterations * args.k_clf)
total_class_loss = total_class_loss / (args.iterations * args.k_clf)
dis_accuracy = total_accuracy / (args.iterations * args.k_disc)
sec_acc = second_acc / (args.iterations * args.k_clf)
clf.feature_extractor = target_model
tar_accuarcy = test(args, clf)
final_accs.append(tar_accuarcy)
if tar_accuarcy > best_tar_acc:
best_tar_acc = tar_accuarcy
torch.save(clf.state_dict(), 'trained_models/adda.pt')
tqdm.write(
f'EPOCH {epoch:03d}: discriminator_loss={mean_loss:.4f}, adv_loss = {mean_adv_loss:.4f}, '
f'discriminator_accuracy={dis_accuracy:.4f}, tar_accuary = {tar_accuarcy:.4f}, best_accuracy = {best_tar_acc:.4f}, '
f'sec_acc = {sec_acc:.4f}, total_class_loss: {total_class_loss:.4f}'
)
# Create the full target model and save it
clf.feature_extractor = target_model
#torch.save(clf.state_dict(), 'trained_models/adda.pt')
jd = {"test_acc": final_accs}
with open(str(args.seed) + '/acc' + str(args.person) + '.json', 'w') as f:
json.dump(jd, f)
def main(args):
# TODO: add DTN model
model = Net().to(device)
model.load_state_dict(torch.load(args.MODEL_FILE))
feature_extractor = model.feature_extractor
clf = model.classifier
discriminator = nn.Sequential(GradientReversal(), nn.Linear(320, 50),
nn.ReLU(), nn.Linear(50, 20), nn.ReLU(),
nn.Linear(20, 1)).to(device)
half_batch = args.batch_size // 2
if args.adapt_setting == 'mnist2mnistm':
source_dataset = MNIST(config.DATA_DIR / 'mnist',
train=True,
download=True,
transform=Compose(
[GrayscaleToRgb(),
ToTensor()]))
target_dataset = MNISTM(train=False)
elif args.adapt_setting == 'svhn2mnist':
source_dataset = ImageClassdata(txt_file=args.src_list,
root_dir=args.src_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
]))
target_dataset = ImageClassdata(txt_file=args.tar_list,
root_dir=args.tar_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.ToTensor(),
]))
elif args.adapt_setting == 'mnist2usps':
source_dataset = ImageClassdata(txt_file=args.src_list,
root_dir=args.src_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.ToTensor(),
]))
target_dataset = ImageClassdata(txt_file=args.tar_list,
root_dir=args.tar_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
]))
else:
raise NotImplementedError
source_loader = DataLoader(source_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True)
target_loader = DataLoader(target_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True)
optim = torch.optim.Adam(
list(discriminator.parameters()) + list(model.parameters()))
if not os.path.exists('logs'): os.makedirs('logs')
f = open(f'logs/{args.adapt_setting}_{args.name}.txt', 'w+')
for epoch in range(1, args.epochs + 1):
batches = zip(source_loader, target_loader)
n_batches = min(len(source_loader), len(target_loader))
total_domain_loss = total_label_accuracy = 0
target_label_accuracy = 0
for (source_x,
source_labels), (target_x,
target_labels) in tqdm(batches,
leave=False,
total=n_batches):
x = torch.cat([source_x, target_x])
x = x.to(device)
domain_y = torch.cat([
torch.ones(source_x.shape[0]),
torch.zeros(target_x.shape[0])
])
domain_y = domain_y.to(device)
label_y = source_labels.to(device)
features = feature_extractor(x).view(x.shape[0], -1)
domain_preds = discriminator(features).squeeze()
label_preds = clf(features[:source_x.shape[0]])
domain_loss = F.binary_cross_entropy_with_logits(
domain_preds, domain_y)
label_loss = F.cross_entropy(label_preds, label_y)
loss = domain_loss + label_loss
optim.zero_grad()
loss.backward()
optim.step()
total_domain_loss += domain_loss.item()
total_label_accuracy += (
label_preds.max(1)[1] == label_y).float().mean().item()
target_label_preds = clf(features[source_x.shape[0]:])
target_label_accuracy += (target_label_preds.cpu().max(1)[1] ==
target_labels).float().mean().item()
mean_loss = total_domain_loss / n_batches
mean_accuracy = total_label_accuracy / n_batches
target_mean_accuracy = target_label_accuracy / n_batches
tqdm.write(
f'EPOCH {epoch:03d}: domain_loss={mean_loss:.4f}, '
f'source_accuracy={mean_accuracy:.4f}, target_accuracy={target_mean_accuracy:.4f}'
)
f.write(
f'EPOCH {epoch:03d}: domain_loss={mean_loss:.4f}, '
f'source_accuracy={mean_accuracy:.4f}, target_accuracy={target_mean_accuracy:.4f}\n'
)
torch.save(
model.state_dict(),
f'trained_models/{args.adapt_setting}_{args.name}_ep{epoch}.pt')
f.close()
class NNet():
"""
Wrapper to manage neural net.
"""
def __init__(self, args):
self.args = args
self.num_channels = NUM_CHANNELS
if args.netType == 1:
self.net = Net(self.num_channels, args)
elif args.netType == 2:
self.net = Net2(self.num_channels, args)
if args.cuda:
self.net = self.net.cuda()
self.load_dataset_from_folder()
self.writer = SummaryWriter()
self.unique_tok = str(time.time())
self.init_weights()
def init_weights(self):
"""
Initialize by Xavier weights
"""
self.net.apply(init_weights)
def load_dataset_from_folder(self):
"""
Load complete dataset
"""
all_data_path = self.args.all_data_path
validation_split_size = self.args.validation_split_size
batch_size = self.args.batch_size
num_workers = self.args.num_workers
shuffle = self.args.shuffle
all_data = ImageFolder(root=all_data_path, transform=TRANSFORM)
classes = all_data.classes
self.classes = classes
validation_size = int(validation_split_size * len(all_data))
test_size = int(validation_split_size * len(all_data))
train_size = len(all_data) - 2 * validation_size
train_dataset, val_dataset, test_dataset = random_split(
all_data, [train_size, validation_size, test_size])
training_data_loader = DataLoader(train_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle)
validation_dataset_loader = DataLoader(val_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle)
test_dataset_loader = DataLoader(test_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle)
self.train_loader = training_data_loader
self.val_loader = validation_dataset_loader
self.test_loader = test_dataset_loader
def train(self):
"""
Train Neural Net
"""
if self.args.optim == 'RMSprop':
optimizer = optim.RMSprop(self.net.parameters(),
lr=self.args.lr,
momentum=self.args.momentum,
weight_decay=self.args.l2_regularization)
elif self.args.optim == 'SGD':
optimizer = optim.SGD(self.net.parameters(),
lr=self.args.lr,
momentum=self.args.momentum)
elif self.args.optim == 'Adam':
optimizer = optim.Adam(self.net.parameters(), lr=self.args.lr)
criterion = nn.CrossEntropyLoss()
# scheduler = optim.lr_scheduler.StepLR(
# optimizer, step_size=self.args.scheduler_step_size, gamma=self.args.scheduler_gamma)
self.net.train()
for epoch in range(self.args.epoch):
start_time = time.time()
running_loss_t = 0.0
num_batches = 0
y_true = []
y_pred = []
# print('Epoch: {} , LR: {}'.format(epoch+1, scheduler.get_lr()))
for data in tqdm(self.train_loader):
inputs, labels = data
labels_cp = labels.clone()
# imshow(torchvision.utils.make_grid(inputs[:,:3,:,:]))
if len(inputs) < 2:
continue
if self.args.cuda:
inputs = inputs.cuda()
labels = labels.cuda()
outputs = self.net(inputs)
loss = criterion(outputs, labels)
_, predicted = torch.max(outputs, 1)
predicted = predicted.cpu()
for i, pred in enumerate(predicted):
y_pred.append(pred)
y_true.append(labels_cp[i])
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss_t += loss.item()
num_batches += 1
end_time = time.time()
train_f1 = f1_score(y_true, y_pred, average='weighted')
# scheduler.step()
self.save(epoch + 1)
self.writer.add_scalar('Loss/train', running_loss_t / num_batches,
epoch + 1)
self.writer.add_scalar('F1/train', train_f1, epoch + 1)
loss_v, val_f1 = self.get_validation_loss(criterion)
self.writer.add_scalar('Loss/val', loss_v, epoch + 1)
self.writer.add_scalar('F1/val', val_f1, epoch + 1)
print(
"Epoch {} Time {:.2f}s Train-Loss {:.3f} Val-Loss {:.3f} Train-F1 {:.3f} Val-F1 {:.3f}"
.format(epoch + 1, end_time - start_time,
running_loss_t / num_batches, loss_v, train_f1,
val_f1))
def get_validation_loss(self, criterion):
"""
Check validation loss
"""
running_loss = 0.0
num_batches = 0
self.net.eval()
y_true = []
y_pred = []
with torch.no_grad():
for data in tqdm(self.val_loader):
images, labels = data
labels_cp = labels.clone()
if self.args.cuda:
images = images.cuda()
labels = labels.cuda()
outputs = self.net(images)
_, predicted = torch.max(outputs, 1)
predicted = predicted.cpu()
for i, pred in enumerate(predicted):
y_pred.append(pred)
y_true.append(labels_cp[i])
loss = criterion(outputs, labels)
running_loss += loss.item()
num_batches += 1
self.net.train()
val_f1 = f1_score(y_true, y_pred, average='weighted')
return running_loss / num_batches, val_f1
def get_test_accuracy(self):
"""
Check overall accuracy of model
"""
y_true = []
y_pred = []
class_correct = list(0. for i in range(4))
class_total = list(0. for i in range(4))
with torch.no_grad():
for data in tqdm(self.test_loader):
images, labels = data
labels_cp = labels.clone()
if self.args.cuda:
images = images.cuda()
labels = labels.cuda()
outputs = self.net(images)
_, predicted = torch.max(outputs, 1)
predicted = predicted.cpu()
for i, pred in enumerate(predicted):
y_pred.append(pred)
y_true.append(labels_cp[i])
c = (predicted == labels_cp).squeeze()
for i in range(min(self.args.batch_size, len(labels_cp))):
label = labels_cp[i]
class_correct[label] += c[i].item()
class_total[label] += 1
print("Test F1: ", f1_score(y_true, y_pred, average='weighted'))
def save(self, epochs, folder_path="../models/"):
"""
Save Model
"""
dict_save = {'params': self.net.state_dict(), 'classes': self.classes}
name = folder_path + self.unique_tok + '_' + str(epochs) + '.model'
torch.save(dict_save, name)
print('Model saved at {}'.format(name))
return name
def load(self, path):
"""
Load a saved model
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dict_load = torch.load(path, map_location=torch.device(device))
self.net.load_state_dict(dict_load['params'])
return dict_load['classes']
def predict(self, inp):
"""
Predict using net
"""
if self.args.cuda:
inp = inp.cuda()
self.net.eval()
with torch.no_grad():
vals = self.net(inp)
print(vals)
_, predicted = torch.max(vals, 1)
predicted = predicted.cpu()
result_class = self.classes[predicted]
return result_class
def train(data_folder, output_folder, es_patience, epochs, TTA):
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
print("使用デバイス：", device)
train = pd.read_csv(data_folder + '/train.csv')
test = pd.read_csv(data_folder + '/test.csv')
arch = EfficientNet.from_pretrained('efficientnet-b1') #モデル
meta_features = list(train.columns)
meta_features.remove('image_name')
meta_features.remove('target')
meta_features.remove('fold')
#パラメータ各種
oof = np.zeros((len(train), 1))
preds = torch.zeros((len(test), 1), dtype=torch.float32, device=device)
skf = KFold(n_splits=5, shuffle=True, random_state=47)
#===========================
# ループスタート
#===========================
for fold, (idxT, idxV) in enumerate(skf.split(np.arange(15)), 1):
print('=' * 20, 'Fold', fold, '=' * 20)
#kfold
train_idx = train.loc[train['fold'].isin(idxT)].index
val_idx = train.loc[train['fold'].isin(idxV)].index
#学習パラメータ
model_path = f'/model_{fold}.pth'
best_val = 0
patience = es_patience
model = Net(arch=arch, n_meta_features=len(meta_features))
model = model.to(device)
#optimizer
optim = torch.optim.Adam(model.parameters(), lr=0.001)
#scheduler
scheduler = ReduceLROnPlateau(optimizer=optim,
mode='max',
patience=1,
verbose=True,
factor=0.2)
criterion = nn.BCEWithLogitsLoss()
trainDataset = MelanomaDataset(
df=train.iloc[train_idx].reset_index(drop=True),
imfolder=data_folder + '/train',
train=True,
transforms=train_transform,
meta_features=meta_features)
valDataset = MelanomaDataset(
df=train.iloc[val_idx].reset_index(drop=True),
imfolder=data_folder + '/train',
train=True,
transforms=test_transform,
meta_features=meta_features)
testDataset = MelanomaDataset(df=test,
imfolder=data_folder + '/test',
train=False,
transforms=test_transform,
meta_features=meta_features)
train_loader = DataLoader(dataset=trainDataset,
batch_size=64,
shuffle=True,
num_workers=2)
val_loader = DataLoader(dataset=valDataset,
batch_size=16,
shuffle=False,
num_workers=2)
test_loader = DataLoader(dataset=testDataset,
batch_size=16,
shuffle=False,
num_workers=2)
#=====================
# epochs
#=====================
for epoch in range(epochs):
start_time = time.time()
correct = 0
epoch_loss = 0
#train_loop
model.train()
for x, y in train_loader:
x[0] = torch.tensor(x[0], device=device, dtype=torch.float32)
x[1] = torch.tensor(x[1], device=device, dtype=torch.float32)
y = torch.tensor(y, device=device, dtype=torch.float32)
optim.zero_grad()
z = model(x)
loss = criterion(z, y.unsqueeze(1))
loss.backward()
optim.step()
pred = torch.round(torch.sigmoid(z))
correct += (pred.cpu() == y.cpu().unsqueeze(1)).sum().item()
epoch_loss += loss.item()
train_acc = correct / len(train_idx)
model.eval()
val_preds = torch.zeros((len(val_idx), 1),
dtype=torch.float32,
device=device)
with torch.no_grad():
#validation_loop
for j, (x_val, y_val) in enumerate(val_loader):
x_val[0] = torch.tensor(x_val[0],
device=device,
dtype=torch.float32)
x_val[1] = torch.tensor(x_val[1],
device=device,
dtype=torch.float32)
y_val = torch.tensor(y_val,
device=device,
dtype=torch.float32)
z_val = model(x_val)
val_pred = torch.sigmoid(z_val)
val_preds[j *
val_loader.batch_size:j * val_loader.batch_size +
x_val[0].shape[0]] = val_pred
val_acc = accuracy_score(train.iloc[val_idx]['target'].values,
torch.round(val_preds.cpu()))
val_roc = roc_auc_score(train.iloc[val_idx]['target'].values,
val_preds.cpu())
print(
'Epoch{:03}: | Loss：{:.3f} | Train acc：{:.3f} | Val acc：{:.3f} | Val roc_auc：{:.3f} | Training time：{}'
.format(
epoch + 1, epoch_loss, train_acc, val_acc, val_roc,
str(
datetime.timedelta(seconds=time.time() -
start_time))[:7]))
scheduler.step(val_roc)
if val_roc >= best_val:
best_val = val_roc
patience = es_patience
torch.save(model, output_folder + model_path)
else:
patience -= 1
if patience == 0:
print('Early stopping. Best Val roc_auc：{:.3f}'.format(
best_val))
break
model = torch.load(output_folder + model_path)
model.eval()
val_preds = torch.zeros((len(val_idx), 1),
dtype=torch.float32,
device=device)
#evaluation loop
with torch.no_grad():
for j, (x_val, y_val) in enumerate(val_loader):
x_val[0] = torch.tensor(x_val[0],
device=device,
dtype=torch.float32)
x_val[1] = torch.tensor(x_val[1],
device=device,
dtype=torch.float32)
y_val = torch.tensor(y_val, device=device, dtype=torch.float32)
z_val = model(x_val)
val_pred = torch.sigmoid(z_val)
val_preds[j * val_loader.batch_size:j * val_loader.batch_size +
x_val[0].shape[0]] = val_pred
oof[val_idx] = val_preds.cpu().numpy()
for _ in range(TTA):
for i, x_test in enumerate(test_loader):
x_test[0] = torch.tensor(x_test[0],
device=device,
dtype=torch.float32)
x_test[1] = torch.tensor(x_test[1],
device=device,
dtype=torch.float32)
z_test = model(x_test)
z_test = torch.sigmoid(z_test)
preds[i *
test_loader.batch_size:i * test_loader.batch_size +
x_test[0].shape[0]] += z_test
preds /= TTA
preds /= skf.n_splits
return preds, oof
test_set = get_test_set(opt.upscale_factor)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
testing_data_loader = DataLoader(dataset=test_set,
num_workers=opt.threads,
batch_size=opt.testBatchSize,
shuffle=False)
print('===> Building model')
model = Net(upscale_factor=opt.upscale_factor).to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
def train(epoch):
epoch_loss = 0
config.logger.info("Training:")
for iteration, batch in enumerate(training_data_loader, 1):
input, target = batch[0].to(device), batch[1].to(device)
optimizer.zero_grad()
loss = criterion(model(input), target)
epoch_loss += loss.item()
loss.backward()
optimizer.step()
class dueling_agent:
def __init__(self, env, args):
# define some important
self.env = env
self.args = args
# trying to define the network
self.net = Net(self.env.action_space.n)
self.target_net = Net(self.env.action_space.n)
# make sure the target net has the same weights as the network
self.target_net.load_state_dict(self.net.state_dict())
if self.args.cuda:
self.net.cuda()
self.target_net.cuda()
# define the optimizer
self.optimizer = torch.optim.Adam(self.net.parameters(),
lr=self.args.lr)
# define the replay memory
self.buffer = replay_memory(self.args.buffer_size)
# define the linear schedule of the exploration
self.exploration_schedule = linear_schedule(int(self.args.total_timesteps * self.args.exploration_fraction), \
self.args.final_ratio, self.args.init_ratio)
# create the folder to save the models
if not os.path.exists(self.args.save_dir):
os.mkdir(self.args.save_dir)
# set the environment folder
self.model_path = os.path.join(self.args.save_dir, self.args.env_name)
if not os.path.exists(self.model_path):
os.mkdir(self.model_path)
# start to do the training
def learn(self):
episode_reward = [0.0]
obs = np.array(self.env.reset())
td_loss = 0
for timestep in range(self.args.total_timesteps):
explore_eps = self.exploration_schedule.get_value(timestep)
with torch.no_grad():
obs_tensor = self._get_tensors(obs)
action_value = self.net(obs_tensor)
# select actions
action = select_actions(action_value, explore_eps)
# excute actions
obs_, reward, done, _ = self.env.step(action)
obs_ = np.array(obs_)
# tryint to append the samples
self.buffer.add(obs, action, reward, obs_, float(done))
obs = obs_
# add the rewards
episode_reward[-1] += reward
if done:
obs = np.array(self.env.reset())
episode_reward.append(0.0)
if timestep > self.args.learning_starts and timestep % self.args.train_freq == 0:
# start to sample the samples from the replay buffer
batch_samples = self.buffer.sample(self.args.batch_size)
td_loss = self._update_network(batch_samples)
if timestep > self.args.learning_starts and timestep % self.args.target_network_update_freq == 0:
# update the target network
self.target_net.load_state_dict(self.net.state_dict())
if len(episode_reward[-101:-1]) == 0:
mean_reward_per_100 = 0
else:
mean_reward_per_100 = np.mean(episode_reward[-101:-1])
num_episode = len(episode_reward) - 1
if done and num_episode % self.args.display_interval == 0:
print('[{}] Frames: {}, Episode: {}, Mean: {:.3f}, Loss: {:.3f}'.format(datetime.now(), timestep, num_episode, \
mean_reward_per_100, td_loss))
torch.save(self.net.state_dict(),
self.model_path + '/model.pt')
# update the network
def _update_network(self, samples):
obses, actions, rewards, obses_next, dones = samples
# convert the data to tensor
obses = self._get_tensors(obses)
actions = torch.tensor(actions, dtype=torch.int64).unsqueeze(-1)
rewards = torch.tensor(rewards, dtype=torch.float32).unsqueeze(-1)
obses_next = self._get_tensors(obses_next)
dones = torch.tensor(1 - dones, dtype=torch.float32).unsqueeze(-1)
# convert into gpu
if self.args.cuda:
actions = actions.cuda()
rewards = rewards.cuda()
dones = dones.cuda()
# calculate the target value
with torch.no_grad():
q_value_temp = self.net(obses_next)
action_max_idx = torch.argmax(q_value_temp, dim=1, keepdim=True)
target_action_value = self.target_net(obses_next)
target_action_max_value = target_action_value.gather(
1, action_max_idx)
target_action_max_value = target_action_max_value.detach()
# target
expected_value = rewards + self.args.gamma * target_action_max_value * dones
# get the real q value
action_value = self.net(obses)
real_value = action_value.gather(1, actions)
loss = (expected_value - real_value).pow(2).mean()
# start to update
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
# get tensors
def _get_tensors(self, obs):
if obs.ndim == 3:
obs = np.transpose(obs, (2, 0, 1))
obs = np.expand_dims(obs, 0)
elif obs.ndim == 4:
obs = np.transpose(obs, (0, 3, 1, 2))
obs = torch.tensor(obs, dtype=torch.float32)
if self.args.cuda:
obs = obs.cuda()
return obs
def main(args):
source_model = Net().to(device)
source_model.load_state_dict(torch.load(args.MODEL_FILE))
source_model.eval()
set_requires_grad(source_model, requires_grad=False)
clf = source_model
source_model = source_model.feature_extractor
target_model = Net().to(device)
target_model.load_state_dict(torch.load(args.MODEL_FILE))
target_model = target_model.feature_extractor
target_clf = clf.classifier
discriminator = nn.Sequential(nn.Linear(320, 50), nn.ReLU(),
nn.Linear(50, 20), nn.ReLU(),
nn.Linear(20, 1)).to(device)
half_batch = args.batch_size // 2
source_dataset = MNIST(config.DATA_DIR / 'mnist',
train=True,
download=True,
transform=Compose([GrayscaleToRgb(),
ToTensor()]))
source_loader = DataLoader(source_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True)
target_dataset = MNISTM(train=False)
target_loader = DataLoader(target_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True)
discriminator_optim = torch.optim.Adam(discriminator.parameters())
target_optim = torch.optim.Adam(target_model.parameters())
criterion = nn.BCEWithLogitsLoss()
for epoch in range(1, args.epochs + 1):
batch_iterator = zip(loop_iterable(source_loader),
loop_iterable(target_loader))
total_loss = 0
total_accuracy = 0
target_label_accuracy = 0
for _ in trange(args.iterations, leave=False):
# Train discriminator
set_requires_grad(target_model, requires_grad=False)
set_requires_grad(discriminator, requires_grad=True)
for _ in range(args.k_disc):
(source_x, _), (target_x, _) = next(batch_iterator)
source_x, target_x = source_x.to(device), target_x.to(device)
source_features = source_model(source_x).view(
source_x.shape[0], -1)
target_features = target_model(target_x).view(
target_x.shape[0], -1)
discriminator_x = torch.cat([source_features, target_features])
discriminator_y = torch.cat([
torch.ones(source_x.shape[0], device=device),
torch.zeros(target_x.shape[0], device=device)
])
preds = discriminator(discriminator_x).squeeze()
loss = criterion(preds, discriminator_y)
discriminator_optim.zero_grad()
loss.backward()
discriminator_optim.step()
total_loss += loss.item()
total_accuracy += ((
preds >
0).long() == discriminator_y.long()).float().mean().item()
# Train classifier
set_requires_grad(target_model, requires_grad=True)
set_requires_grad(discriminator, requires_grad=False)
for _ in range(args.k_clf):
_, (target_x, target_labels) = next(batch_iterator)
target_x = target_x.to(device)
target_features = target_model(target_x).view(
target_x.shape[0], -1)
# flipped labels
discriminator_y = torch.ones(target_x.shape[0], device=device)
preds = discriminator(target_features).squeeze()
loss = criterion(preds, discriminator_y)
target_optim.zero_grad()
loss.backward()
target_optim.step()
target_label_preds = target_clf(target_features)
target_label_accuracy += (target_label_preds.cpu().max(1)[1] ==
target_labels).float().mean().item()
mean_loss = total_loss / (args.iterations * args.k_disc)
mean_accuracy = total_accuracy / (args.iterations * args.k_disc)
target_mean_accuracy = target_label_accuracy / (args.iterations *
args.k_clf)
tqdm.write(
f'EPOCH {epoch:03d}: discriminator_loss={mean_loss:.4f}, '
f'discriminator_accuracy={mean_accuracy:.4f}, target_accuracy={target_mean_accuracy:.4f}'
)
# Create the full target model and save it
clf.feature_extractor = target_model
torch.save(clf.state_dict(), 'trained_models/adda.pt')
def main(args):
trainloader, testloader, classes = load_dataset(args.dataset)
print(len(trainloader), len(testloader))
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
writer_train = SummaryWriter(log_dir=os.path.join(args.outdir, "train/"), purge_step=0)
writer_test = SummaryWriter(log_dir=os.path.join(args.outdir, "test/"), purge_step=0)
correct = 0
total = 0
iter = 0
for epoch in range(args.epochlen): # loop over the dataset multiple times
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if iter % 2000 == 1999: # print every 2000 mini-batches
print('[%d, %5d] loss: %.3f' % (epoch + 1, iter, running_loss / 2000))
writer_train.add_scalar('Training Loss', running_loss, iter)
running_loss = 0.0
iter += 1
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
train_accuracy = (100 * correct) / total
writer_train.add_scalar('Accuracy', train_accuracy, epoch+1)
test_loss, test_accuracy = test_data(net, classes, testloader, criterion)
print(f"Accuracy of the network on the 10000 test images for epoch {epoch+1} is {test_accuracy}")
writer_test.add_scalar('Testing Loss', test_loss, epoch+1)
writer_test.add_scalar('Accuracy', test_accuracy, epoch+1)
check_class_performance(net, classes, testloader)
# Plot confusion matrix
y_true = []
y_pred = []
for data in testloader:
images, labels = data
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
y_true += labels.tolist()
y_pred += predicted.tolist()
cm = confusion_matrix(np.array(y_true), np.array(y_pred))
print(cm)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=classes)
disp.plot(include_values=True, cmap='viridis', ax=None, xticks_rotation='horizontal', values_format=None)
print("Accuracy score : ", accuracy_score(y_true, y_pred))
print("classification report : ", classification_report(y_true, y_pred))
plt.show()
print('Finished Training')
def train(options):
exp_name = options['exp_name']
batch_size = options['batch_size']
use_pca = options['use_pca']
model_type = options['model_type']
loss_fn = options['loss_fn']
optim = options['optim']
use_scheduler = options['use_scheduler']
lr = options['lr']
epochs = options['epochs']
pca_var_hold = options['pca_var_hold']
debug_mode = options['debug_mode']
if os.path.exists(exp_name):
shutil.rmtree(exp_name)
time.sleep(1)
writer = SummaryWriter(exp_name,flush_secs=1)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
X = os.listdir('hilbert_data')
X_train = X[:int(0.8*len(X))]
X_test = X[int(0.8*len(X)):]
# X = np.load('bined_x.npy')
# y = np.load('bined_y.npy')
# X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# if use_pca and 'Raw' in exp_name:
# scaler = PCA(pca_var_hold)
# scaler.fit(X_train)
# X_train = scaler.transform(X_train)
# X_test = scaler.transform(X_test)
# needed_dim = X_train.shape[1]
dataset_train = HIL_MOOD(X_train, model_type=model_type,data_type='train',debug_mode=debug_mode)
train_loader = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True)
dataset_val = HIL_MOOD(X_test, model_type=model_type,data_type='val')
valid_loader = DataLoader(dataset=dataset_val, batch_size=batch_size, shuffle=False)
model = Net()
model.to(device)
if optim == None:
print('you need to specify an optimizer')
exit()
elif optim == 'adam':
optimizer = torch.optim.Adam( model.parameters(), lr=lr)
elif optim == 'sgd':
optimizer = torch.optim.SGD( model.parameters(), lr=lr,momentum=0.9)
if use_scheduler:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min',verbose=True,threshold=0.0001,patience = 10)
if loss_fn == None:
print('you need to specify an optimizer')
exit()
else:
if loss_fn == 'mse':
loss_fn = torch.nn.MSELoss()
elif loss_fn == 'cross_entropy':
loss_fn = torch.nn.CrossEntropyLoss()
mean_train_losses = []
mean_valid_losses = []
valid_acc_list = []
best = 0 #small number for acc big number for loss to save a model
for epoch in range(epochs):
model.train()
train_losses = []
valid_losses = []
for i, (images, labels) in enumerate(train_loader):
if images.shape[0] != batch_size:
continue
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad()
# print(images.shape)
outputs = model(images)
# print(images.shape)
# print(outputs.shape)
# print(labels.shape)
# print(i)
loss =loss_fn(outputs,labels)
# print('loss: ',loss.item())
writer.add_scalar('Loss/train', loss.item(), len(train_loader)*epoch+i)
loss.backward()
optimizer.step()
train_losses.append(loss.item())
del outputs
# if (i * batch_size) % (batch_size * 100) == 0:
# print(f'{i * batch_size} / 50000')
model.eval()
correct_5_2 = 0
correct_5_1 = 0
total_loss = 0
total = 0
accsat =[0.5,0.05,0.005]
accs = np.zeros(len(accsat))
# corrs = np.zeros(len(accsat))
correct_array = np.zeros(len(accsat))
with torch.no_grad():
for i, (images, labels) in enumerate(valid_loader):
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
loss = loss_fn(outputs, labels)
for i in range(len(accsat)):
correct_array[i] += accat(outputs,labels,thresh=accsat[i])
# total_loss += loss.item()
total += labels.size(0)
valid_losses.append(loss.item())
mean_train_losses.append(np.mean(train_losses))
mean_valid_losses.append(np.mean(valid_losses))
# scheduler.step(np.mean(valid_losses))
for i in range(len(accsat)):
accs[i] = 100*correct_array[i]/total
writer.add_scalar('Acc/val_@'+str(accsat[i]), accs[i], epoch)
if np.mean(valid_losses) < best:
best = np.mean(valid_losses)
torch.save(model.state_dict(),os.path.join(os.getcwd(),'models','meh.pth'))
writer.add_scalar('Loss/val', np.mean(valid_losses), epoch)
# valid_acc_list.append(accuracy)
if epoch ==epochs-1:
print('epoch : {}, train loss : {:.4f}, valid loss : {:.4f}, [email protected] : {:.4f}'\
.format(epoch+1, np.mean(train_losses), np.mean(valid_losses), accsat[1]))
class TrainBigramNN(tune.Trainable):
def _setup(self, config):
print("Loading word vectors...")
word2index, word_vecs = process_word_vecs(FAST_TEXT)
# Note that the word embeddings are normalized.
self.wv = WV(F.normalize(word_vecs), word2index)
# wv = WV(word_vecs, word2index)
print("Done.")
self.corpus_size = config["corpus_size"]
bigram_fn_name = "diff"
out_bigram_dim = 300
dist_fn_name = "cos_dist"
loss_fn_name = "mrl"
margin = config["margin"]
self.lr = config["lr"]
self.num_epochs = config["num_epochs"]
self.batch_size = config["batch_size"]
self.test_model = True
self.test_freq = config["test_freq"]
with open(PROCESSED / "train.{}.pkl".format(str(self.corpus_size)), "rb") as f:
wiki_train = pickle.load(f)
with open(PROCESSED / "valid.pkl", "rb") as f:
wiki_valid = pickle.load(f)
wiki_combined = wiki_train + wiki_valid
self.corpus = Corpus("wiki", wiki_combined, self.wv, filter_stopwords=True)
self.model = Net(
self.wv.vecs.size(1), BigramEncoder(bigram_fn_name), out_bigram_dim
)
self.model.to(device)
self.dist_fn = DistanceFunction(dist_fn_name)
self.loss_fn = LossFunction(loss_fn_name, margin=margin)
self.device = device
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
torch.manual_seed(config["seed"])
print("Traninig on Wikipedia corpus of size {}".format(self.corpus_size))
def _train(self):
result = train(
self.wv,
self.corpus.ix_sents[: self.corpus_size],
self.corpus.sent_lengths[: self.corpus_size],
self.corpus.ix_sents[self.corpus_size :],
self.corpus.sent_lengths[self.corpus_size :],
self.model,
self.wv.vecs,
self.dist_fn,
self.loss_fn,
self.optimizer,
self.lr,
self.num_epochs,
self.batch_size,
self._iteration,
self.test_model,
self.test_freq,
self.device,
)
return result
def _save(self, tmp_checkpoint_dir):
checkpoint_path = str(Path(tmp_checkpoint_dir) / "model.pth")
torch.save(self.model.state_dict(), checkpoint_path)
return checkpoint_path
def _restore(self, tmp_checkpoint_dir):
checkpoint_path = str(Path(tmp_checkpoint_dir) / "model.pth")
self.model.load_state_dict(torch.load(checkpoint_path))
def main():
check_dir = '../LPSfiles/' + name
if not os.path.exists(check_dir):
os.mkdir(check_dir)
# data
val_loader = torch.utils.data.DataLoader(PriorFolder(opt.val_dir,
opt.prior_dir,
size=256,
mean=mean,
std=std),
batch_size=opt.b * 3,
shuffle=False,
num_workers=4,
pin_memory=True)
train_loader = torch.utils.data.DataLoader(Folder(opt.train_dir,
scales=[64] * 3 +
[128, 256],
crop=0.9,
flip=True,
rotate=None,
mean=mean,
std=std),
collate_fn=collate_more,
batch_size=opt.b * 6,
shuffle=True,
num_workers=4,
pin_memory=True)
# models
p = 5
net = Net(base=opt.base)
fcn = FCN(net)
net = nn.DataParallel(net).cuda()
net.train()
"""
# fcn = nn.DataParallel(fcn).cuda()
# sdict =torch.load('/home/crow/LPSfiles/Train2_vgg16/fcn-iter13800.pth')
# fcn.load_state_dict(sdict)
fcn = nn.DataParallel(fcn).cuda()
fcn.train()
optimizer = torch.optim.Adam([
{'params': fcn.parameters(), 'lr': 1e-4},
])
logs = {'best_it':0, 'best': 0}
sal_data_iter = iter(train_loader)
i_sal_data = 0
for it in tqdm(range(opt.max)):
# for it in tqdm(range(1)):
# if it > 1000 and it % 100 == 0:
# optimizer.param_groups[0]['lr'] *= 0.5
if i_sal_data >= len(train_loader):
sal_data_iter = iter(train_loader)
i_sal_data = 0
data, lbls, _ = sal_data_iter.next()
i_sal_data += 1
data = data.cuda()
lbls = [lbl.unsqueeze(1).cuda() for lbl in lbls]
msks = fcn(data)
loss = sum([F.binary_cross_entropy_with_logits(msk, lbl) for msk, lbl in zip(msks, lbls)])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if it % 10 == 0:
writer.add_scalar('loss', loss.item(), it)
image = make_image_grid(data[:6], mean, std)
writer.add_image('Image', torchvision.utils.make_grid(image), it)
big_msk = F.sigmoid(msks[-1]).expand(-1, 3, -1, -1)
writer.add_image('msk', torchvision.utils.make_grid(big_msk.data[:6]), it)
big_msk = lbls[-1].expand(-1, 3, -1, -1)
writer.add_image('gt', torchvision.utils.make_grid(big_msk.data[:6]), it)
# if it % 100 == 0:
if it != 0 and it % 100 == 0:
fm, mae = validate(val_loader, fcn, os.path.join(check_dir, 'results'),
os.path.join(opt.val_dir, 'masks'))
print(u'损失: %.4f'%(loss.item()))
print(u'最大FM: iteration %d的%.4f, 这次FM: %.4f'%(logs['best_it'], logs['best'], fm))
logs[it] = {'FM': fm}
if fm > logs['best']:
logs['best'] = fm
logs['best_it'] = it
torch.save(fcn.state_dict(), '%s/fcn-best.pth' % (check_dir))
with open(os.path.join(check_dir, 'logs.json'), 'w') as outfile:
json.dump(logs, outfile)
torch.save(fcn.state_dict(), '%s/fcn-iter%d.pth' % (check_dir, it))
"""
###################################################################################################
val_loader = torch.utils.data.DataLoader(PriorFolder(opt.val_dir,
opt.prior_dir,
size=256,
mean=mean,
std=std),
batch_size=opt.b,
shuffle=False,
num_workers=4,
pin_memory=True)
train_loader = torch.utils.data.DataLoader(Folder(opt.train_dir,
scales=[256],
crop=0.9,
flip=True,
rotate=None,
mean=mean,
std=std),
collate_fn=collate_more,
batch_size=opt.b,
shuffle=True,
num_workers=4,
pin_memory=True)
optimizer = torch.optim.Adam([
{
'params': net.parameters(),
'lr': 1e-4
},
])
logs = {'best_it': 0, 'best': 0}
sal_data_iter = iter(train_loader)
i_sal_data = 0
for it in tqdm(range(opt.max)):
# if it > 1000 and it % 100 == 0:
# optimizer.param_groups[0]['lr'] *= 0.5
if i_sal_data >= len(train_loader):
sal_data_iter = iter(train_loader)
i_sal_data = 0
data, lbl, _ = sal_data_iter.next()
i_sal_data += 1
data = data.cuda()
lbl = lbl[0].unsqueeze(1)
noisy_label = (lbl.numpy() +
np.random.binomial(1,
float(p) / 100.0, (256, 256))) % 2
noisy_label = torch.Tensor(noisy_label).cuda()
lbl = lbl.cuda()
msk = net(data, noisy_label)
loss = F.binary_cross_entropy_with_logits(msk, lbl)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if it % 10 == 0:
writer.add_scalar('loss', loss.item(), it)
image = make_image_grid(data[:6], mean, std)
writer.add_image('Image', torchvision.utils.make_grid(image), it)
big_msk = F.sigmoid(msk).expand(-1, 3, -1, -1)
writer.add_image('msk',
torchvision.utils.make_grid(big_msk.data[:6]), it)
big_msk = lbl.expand(-1, 3, -1, -1)
writer.add_image('gt',
torchvision.utils.make_grid(big_msk.data[:6]), it)
# if it % 200 == 0:
if it != 0 and it % 100 == 0:
fm, mae = validate(val_loader, net,
os.path.join(check_dir, 'results'),
os.path.join(opt.val_dir, 'masks'))
print(u'损失: %.4f' % (loss.item()))
print(u'最大FM: iteration %d的%.4f, 这次FM: %.4f' %
(logs['best_it'], logs['best'], fm))
logs[it] = {'FM': fm}
if fm > logs['best']:
logs['best'] = fm
logs['best_it'] = it
torch.save(net.state_dict(), '%s/net-best.pth' % (check_dir))
with open(os.path.join(check_dir, 'logs.json'), 'w') as outfile:
json.dump(logs, outfile)
torch.save(net.state_dict(), '%s/net-iter%d.pth' % (check_dir, it))
data_tools = DataTools()
transformed_training_dataset = data_tools.transformed_training_data
# load training data in batches
batch_size = args.batch_size
num_workers = args.num_workers
train_loader = DataLoader(transformed_training_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
# Define the loss and optimization
lr = args.lr
criterion = nn.SmoothL1Loss()
optimizer = optim.Adam(net.parameters(), lr=lr)
def train_net(n_epochs):
# prepare the net for training
net.train()
for epoch in range(n_epochs): # loop over the dataset multiple times
running_loss = 0.0
# train on batches of data, assumes you already have train_loader
for batch_i, data in enumerate(train_loader):
# get the input images and their corresponding labels
images = data['image']
from prepare_data import loader_tr, loader_ts
import myloss
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net(50, 50, 50)
#criterion = nn.MSELoss()
criterion = nn.L1Loss(reduction='mean')
#criterion = nn.KLDivLoss(reduction='batchmean')
#criterion = nn.HingeEmbeddingLoss(reduction='sum')
#criterion = nn.CosineEmbeddingLoss()
#criterion = myloss.HingeLoss()
lr = 0.01
momentum = 0.9
#optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
#optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=5e-4)
model.to(device)
losses_tr = []
losses_ts = []
num_epoches = 100
#训练
for epoch in range(num_epoches):
print(f'epoch {epoch}, training....')
loss_tr = 0
acc_tr = 0
if epoch % 5 == 0:
#每隔5轮
test_dataset = FacialKeypointsDataset(
csv_file='data/test_frames_keypoints.csv',
root_dir='data/test/',
transform=data_transform)
# load test data in batches
batch_size = 10
test_loader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=4)
## TODO: Define the loss and optimization
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
"""
with learning rate of 0.001, β1=0.9, β2=0.999 and " of 1e− 08, is used for minimizing Mean Squared Error (MSE).
"""
def train_net(n_epochs):
# prepare the net for training
net.train()
for epoch in range(n_epochs): # loop over the dataset multiple times
running_loss = 0.0
# train on batches of data, assumes you already have train_loader
for batch_i, data in enumerate(train_loader):
**kwargs)
test_loader = torch.utils.data.DataLoader(datasets.MNIST(
'datas',
train=False,
transform=transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
model = Net()
if args.cuda:
model.cuda() # 将所有的模型参数移动到GPU上
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
def train(epoch):
model.train() # 把module设成training模式，对Dropout和BatchNorm有影响
for batch_idx, (data, target) in enumerate(train_loader):
if args.cuda:
data, target = data.cuda(), target.cuda()
# Variable类对Tensor对象进行封装，会保存该张量对应的梯度，以及对生成该张量的函数grad_fn的一个引用。如果该张量是用户创建的，grad_fn是None，称这样的Variable为叶子Variable。
data, target = Variable(data), Variable(target)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target) # 负log似然损失
loss.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
def main(args):
model = Net().to(device)
model.load_state_dict(torch.load(args.MODEL_FILE))
feature_extractor = model.feature_extractor
clf = model.classifier
discriminator = nn.Sequential(GradientReversal(), nn.Linear(320, 50),
nn.ReLU(), nn.Linear(50, 20), nn.ReLU(),
nn.Linear(20, 1)).to(device)
half_batch = args.batch_size // 2
source_dataset = MNIST(config.DATA_DIR / 'mnist',
train=True,
download=True,
transform=Compose([GrayscaleToRgb(),
ToTensor()]))
source_loader = DataLoader(source_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True)
target_dataset = MNISTM(train=False)
target_loader = DataLoader(target_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True)
optim = torch.optim.Adam(
list(discriminator.parameters()) + list(model.parameters()))
for epoch in range(1, args.epochs + 1):
batches = zip(source_loader, target_loader)
n_batches = min(len(source_loader), len(target_loader))
total_domain_loss = total_label_accuracy = 0
target_label_accuracy = 0
for (source_x,
source_labels), (target_x,
target_labels) in tqdm(batches,
leave=False,
total=n_batches):
x = torch.cat([source_x, target_x])
x = x.to(device)
domain_y = torch.cat([
torch.ones(source_x.shape[0]),
torch.zeros(target_x.shape[0])
])
domain_y = domain_y.to(device)
label_y = source_labels.to(device)
features = feature_extractor(x).view(x.shape[0], -1)
domain_preds = discriminator(features).squeeze()
label_preds = clf(features[:source_x.shape[0]])
domain_loss = F.binary_cross_entropy_with_logits(
domain_preds, domain_y)
label_loss = F.cross_entropy(label_preds, label_y)
loss = domain_loss + label_loss
optim.zero_grad()
loss.backward()
optim.step()
total_domain_loss += domain_loss.item()
total_label_accuracy += (
label_preds.max(1)[1] == label_y).float().mean().item()
target_label_preds = clf(features[source_x.shape[0]:])
target_label_accuracy += (target_label_preds.cpu().max(1)[1] ==
target_labels).float().mean().item()
mean_loss = total_domain_loss / n_batches
mean_accuracy = total_label_accuracy / n_batches
target_mean_accuracy = target_label_accuracy / n_batches
tqdm.write(
f'EPOCH {epoch:03d}: domain_loss={mean_loss:.4f}, '
f'source_accuracy={mean_accuracy:.4f}, target_accuracy={target_mean_accuracy:.4f}'
)
torch.save(model.state_dict(), 'trained_models/revgrad.pt')
def main(args):
clf_model = Net().to(device)
clf_model.load_state_dict(torch.load(args.MODEL_FILE))
feature_extractor = clf_model.feature_extractor
discriminator = clf_model.classifier
critic = nn.Sequential(nn.Linear(320, 50), nn.ReLU(), nn.Linear(50, 20),
nn.ReLU(), nn.Linear(20, 1)).to(device)
half_batch = args.batch_size // 2
source_dataset = MNIST(config.DATA_DIR / 'mnist',
train=True,
download=True,
transform=Compose([GrayscaleToRgb(),
ToTensor()]))
source_loader = DataLoader(source_dataset,
batch_size=half_batch,
drop_last=True,
shuffle=True,
num_workers=0,
pin_memory=True)
target_dataset = MNISTM(train=False)
target_loader = DataLoader(target_dataset,
batch_size=half_batch,
drop_last=True,
shuffle=True,
num_workers=0,
pin_memory=True)
critic_optim = torch.optim.Adam(critic.parameters(), lr=1e-4)
clf_optim = torch.optim.Adam(clf_model.parameters(), lr=1e-4)
clf_criterion = nn.CrossEntropyLoss()
for epoch in range(1, args.epochs + 1):
batch_iterator = zip(loop_iterable(source_loader),
loop_iterable(target_loader))
total_loss = 0
total_accuracy = 0
for _ in trange(args.iterations, leave=False):
(source_x, source_y), (target_x, _) = next(batch_iterator)
# Train critic
set_requires_grad(feature_extractor, requires_grad=False)
set_requires_grad(critic, requires_grad=True)
source_x, target_x = source_x.to(device), target_x.to(device)
source_y = source_y.to(device)
with torch.no_grad():
h_s = feature_extractor(source_x).data.view(
source_x.shape[0], -1)
h_t = feature_extractor(target_x).data.view(
target_x.shape[0], -1)
for _ in range(args.k_critic):
gp = gradient_penalty(critic, h_s, h_t)
critic_s = critic(h_s)
critic_t = critic(h_t)
wasserstein_distance = critic_s.mean() - critic_t.mean()
critic_cost = -wasserstein_distance + args.gamma * gp
critic_optim.zero_grad()
critic_cost.backward()
critic_optim.step()
total_loss += critic_cost.item()
# Train classifier
set_requires_grad(feature_extractor, requires_grad=True)
set_requires_grad(critic, requires_grad=False)
for _ in range(args.k_clf):
source_features = feature_extractor(source_x).view(
source_x.shape[0], -1)
target_features = feature_extractor(target_x).view(
target_x.shape[0], -1)
source_preds = discriminator(source_features)
clf_loss = clf_criterion(source_preds, source_y)
wasserstein_distance = critic(source_features).mean() - critic(
target_features).mean()
loss = clf_loss + args.wd_clf * wasserstein_distance
clf_optim.zero_grad()
loss.backward()
clf_optim.step()
mean_loss = total_loss / (args.iterations * args.k_critic)
tqdm.write(f'EPOCH {epoch:03d}: critic_loss={mean_loss:.4f}')
torch.save(clf_model.state_dict(), 'trained_models/wdgrl.pt')
class a2c_agent:
def __init__(self, envs, args):
self.envs = envs
self.args = args
# define the network
self.net = Net(self.envs.action_space.n)
if self.args.cuda:
self.net.cuda()
# define the optimizer
self.optimizer = torch.optim.RMSprop(self.net.parameters(),
lr=self.args.lr,
eps=self.args.eps,
alpha=self.args.alpha)
if not os.path.exists(self.args.save_dir):
os.mkdir(self.args.save_dir)
# check the saved path for envs..
self.model_path = self.args.save_dir + self.args.env_name + '/'
if not os.path.exists(self.model_path):
os.mkdir(self.model_path)
# get the obs..
self.batch_ob_shape = (self.args.num_processes * self.args.nsteps,
) + self.envs.observation_space.shape
self.obs = np.zeros(
(self.args.num_processes, ) + self.envs.observation_space.shape,
dtype=self.envs.observation_space.dtype.name)
self.obs[:] = self.envs.reset()
self.dones = [False for _ in range(self.args.num_processes)]
# train the network..
def learn(self):
if not self.args.no_sil:
sil_model = sil_module(self.net, self.args, self.optimizer)
num_updates = self.args.total_frames // (self.args.num_processes *
self.args.nsteps)
# get the reward to calculate other information
episode_rewards = torch.zeros([self.args.num_processes, 1])
final_rewards = torch.zeros([self.args.num_processes, 1])
# start to update
for update in range(num_updates):
mb_obs, mb_rewards, mb_actions, mb_dones = [], [], [], []
for step in range(self.args.nsteps):
with torch.no_grad():
input_tensor = self._get_tensors(self.obs)
_, pi = self.net(input_tensor)
# select actions
actions = select_actions(pi)
cpu_actions = actions.squeeze(1).cpu().numpy()
# start to store the information
mb_obs.append(np.copy(self.obs))
mb_actions.append(cpu_actions)
mb_dones.append(self.dones)
# step
obs, rewards, dones, _ = self.envs.step(cpu_actions)
# process rewards...
raw_rewards = copy.deepcopy(rewards)
rewards = np.sign(rewards)
# start to store the rewards
self.dones = dones
if not self.args.no_sil:
sil_model.step(input_tensor.detach().cpu().numpy(),
cpu_actions, raw_rewards, dones)
mb_rewards.append(rewards)
for n, done in enumerate(dones):
if done:
self.obs[n] = self.obs[n] * 0
self.obs = obs
raw_rewards = torch.from_numpy(
np.expand_dims(np.stack(raw_rewards), 1)).float()
episode_rewards += raw_rewards
# get the masks
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in dones])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
# update the obs
mb_dones.append(self.dones)
# process the rollouts
mb_obs = np.asarray(mb_obs, dtype=np.uint8).swapaxes(1, 0).reshape(
self.batch_ob_shape)
mb_rewards = np.asarray(mb_rewards,
dtype=np.float32).swapaxes(1, 0)
mb_actions = np.asarray(mb_actions, dtype=np.int32).swapaxes(1, 0)
mb_dones = np.asarray(mb_dones, dtype=np.bool).swapaxes(1, 0)
mb_masks = mb_dones[:, :-1]
mb_dones = mb_dones[:, 1:]
with torch.no_grad():
input_tensor = self._get_tensors(self.obs)
last_values, _ = self.net(input_tensor)
# compute returns
for n, (rewards, dones, value) in enumerate(
zip(mb_rewards, mb_dones,
last_values.detach().cpu().numpy().squeeze())):
rewards = rewards.tolist()
dones = dones.tolist()
if dones[-1] == 0:
rewards = discount_with_dones(rewards + [value],
dones + [0],
self.args.gamma)[:-1]
else:
rewards = discount_with_dones(rewards, dones,
self.args.gamma)
mb_rewards[n] = rewards
mb_rewards = mb_rewards.flatten()
mb_actions = mb_actions.flatten()
# start to update network
vl, al, ent = self._update_network(mb_obs, mb_rewards, mb_actions)
# start to update the sil_module
if not self.args.no_sil:
mean_adv, num_samples = sil_model.train_sil_model()
if update % self.args.log_interval == 0:
if not self.args.no_sil:
print('[{}] Update: {}/{}, Frames: {}, Rewards: {:.2f}, VL: {:.3f}, PL: {:.3f},' \
'Ent: {:.2f}, Min: {}, Max:{}, BR:{}, E:{}, VS:{}, S:{}'.format(\
datetime.now(), update, num_updates, (update+1)*(self.args.num_processes * self.args.nsteps),\
final_rewards.mean(), vl, al, ent, final_rewards.min(), final_rewards.max(), sil_model.get_best_reward(), \
sil_model.num_episodes(), num_samples, sil_model.num_steps()))
else:
print('[{}] Update: {}/{}, Frames: {}, Rewards: {:.2f}, VL: {:.3f}, PL: {:.3f},' \
'Ent: {:.2f}, Min: {}, Max:{}'.format(\
datetime.now(), update, num_updates, (update+1)*(self.args.num_processes * self.args.nsteps),\
final_rewards.mean(), vl, al, ent, final_rewards.min(), final_rewards.max()))
torch.save(self.net.state_dict(), self.model_path + 'model.pt')
# update_network
def _update_network(self, obs, returns, actions):
# evaluate the actions
input_tensor = self._get_tensors(obs)
values, pi = self.net(input_tensor)
# define the tensor of actions, returns
returns = torch.tensor(returns, dtype=torch.float32).unsqueeze(1)
actions = torch.tensor(actions, dtype=torch.int64).unsqueeze(1)
if self.args.cuda:
returns = returns.cuda()
actions = actions.cuda()
# evaluate actions
action_log_probs, dist_entropy = evaluate_actions(pi, actions)
# calculate advantages...
advantages = returns - values
# get the value loss
value_loss = advantages.pow(2).mean()
# get the action loss
action_loss = -(advantages.detach() * action_log_probs).mean()
# total loss
total_loss = action_loss + self.args.value_loss_coef * value_loss - self.args.entropy_coef * dist_entropy
# start to update
self.optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(self.net.parameters(),
self.args.max_grad_norm)
self.optimizer.step()
return value_loss.item(), action_loss.item(), dist_entropy.item()
# get the tensors...
def _get_tensors(self, obs):
input_tensor = torch.tensor(np.transpose(obs, (0, 3, 1, 2)),
dtype=torch.float32)
if self.args.cuda:
input_tensor = input_tensor.cuda()
return input_tensor
# call show_all_keypoints
show_all_keypoints(np.squeeze(image), predicted_key_pts,
ground_truth_pts)
plt.axis('off')
plt.show()
# visualize_output(test_images, test_outputs, gt_pts)
import torch.optim as optim
# Define the loss and optimization
criterion = nn.SmoothL1Loss()
optimizer = optim.Adam(net.parameters(), lr=0.0005)
def train_net(n_epochs):
# prepare the net for training
net.train()
for epoch in range(n_epochs): # loop over the dataset multiple times
running_loss = 0.0
# train on batches of data, assumes you already have train_loader
for batch_i, data in enumerate(train_loader):
# get the input images and their corresponding labels
images = data['image']
