emilyxuan/accesscontrol
0
1import os2import warnings3import datetime4import csv5import json6import random7 8from absl import app, flags9import torch10from tensorboardX import SummaryWriter11from torchvision.utils import make_grid, save_image12from tqdm import tqdm13 14from clients import ClientsGroupMultiTargetAttackedNonIID15from model import UNet16from trigger_encoder import TriggerEncoder17 18FLAGS = flags.FLAGS19 20# ===== Basic =====21flags.DEFINE_string('logdir', None, help='log directory')22flags.DEFINE_bool('train', False, help='run training')23flags.DEFINE_string('dataset_name', 'cifar10', help='dataset name: cifar10 / celeba')24flags.DEFINE_string('device', 'auto', help='auto / cpu / cuda / cuda:0 ...')25flags.DEFINE_integer('seed', 42, help='global random seed')26flags.DEFINE_string('resume_ckpt', None, help='path to resume checkpoint')27flags.DEFINE_string('base_logdir', './logs', help='base log directory')28flags.DEFINE_string('run_name', '', help='optional fixed run name; empty means auto timestamp name')29 30# ===== UNet =====31flags.DEFINE_integer('ch', 128, help='base channel of UNet')32flags.DEFINE_multi_integer('ch_mult', [1, 2, 2, 2], help='channel multiplier')33flags.DEFINE_multi_integer('attn', [1], help='add attention to these levels')34flags.DEFINE_integer('num_res_blocks', 2, help='# resblock in each level')35flags.DEFINE_float('dropout', 0.1, help='dropout rate of resblock')36flags.DEFINE_integer('cond_dim', 128, help='condition embedding dimension')37 38# ===== Trigger encoder / conditional training =====39flags.DEFINE_float('neg_ratio', 0.2, help='negative sample ratio inside each batch')40flags.DEFINE_integer('trigger_base_ch', 32, help='base channel of trigger encoder')41flags.DEFINE_integer('unseen_num_points', 2, help='number of active points in unseen trigger')42flags.DEFINE_integer('unseen_grid_size', 8, help='grid size for unseen trigger sampling')43flags.DEFINE_integer('unseen_square_size', 3, help='square size for unseen trigger blocks')44flags.DEFINE_integer('unseen_base_seed', 12345, help='base seed for unseen trigger generation')45flags.DEFINE_integer('blur_kernel_size', 5, help='kernel size for blurred negative images')46flags.DEFINE_integer('blur_repeat', 2, help='how many times to apply avg-pool blur')47 48# ===== Gaussian Diffusion =====49flags.DEFINE_float('beta_1', 1e-4, help='start beta value')50flags.DEFINE_float('beta_T', 0.02, help='end beta value')51flags.DEFINE_integer('T', 1000, help='total diffusion steps')52 53# ===== Training =====54flags.DEFINE_float('lr', 2e-4, help='target learning rate')55flags.DEFINE_integer('img_size', 32, help='image size; overwritten by dataset defaults if needed')56flags.DEFINE_integer('batch_size', 128, help='batch size')57flags.DEFINE_integer('num_workers', 4, help='workers of Dataloader')58flags.DEFINE_bool('parallel', False, help='multi gpu training')59flags.DEFINE_string('split_mode', 'iid', help='iid or noniid')60 61# ===== Fed =====62flags.DEFINE_string('mode', 'access_control', help='training mode: unconditional / conditional_only / access_control')63flags.DEFINE_integer('mid_T', 500, help='mid T split local global')64flags.DEFINE_integer('local_epoch', 1, help='local epoch')65flags.DEFINE_integer('total_round', 300, help='total round')66flags.DEFINE_integer('client_num', 5, help='client num')67flags.DEFINE_integer('save_round', 20, help='save full checkpoint every N rounds; 0 disables')68flags.DEFINE_integer('data_distribution_seed', 42, help='data distribution seed')69flags.DEFINE_float('ema_scale', 0.9999, help='EMA scale')70 71# ===== Sampling / Logging =====72flags.DEFINE_integer('sample_batch_size', 4, help='number of x_T samples for visual sampling')73flags.DEFINE_integer('sample_every', 1, help='save sample images every N rounds; 0 disables')74flags.DEFINE_bool('save_samples', True, help='whether to save sample grids')75flags.DEFINE_bool('save_tensorboard', True, help='whether to write tensorboard images')76flags.DEFINE_bool('save_client_states', True, help='whether to save per-client states inside checkpoint; disable to reduce checkpoint size')77flags.DEFINE_bool('save_global_model', True, help='save raw global model in checkpoint')78flags.DEFINE_bool('save_global_ema_model', True, help='save global EMA model in checkpoint')79flags.DEFINE_integer('sample_start_step', 0, help='sampling start diffusion step')80flags.DEFINE_integer('sample_end_step', 1000, help='sampling end diffusion step')81 82def get_device() -> torch.device:83 if FLAGS.device == 'auto':84 return torch.device('cuda' if torch.cuda.is_available() else 'cpu')85 return torch.device(FLAGS.device)86 87 88def set_seed(seed: int) -> None:89 random.seed(seed)90 torch.manual_seed(seed)91 if torch.cuda.is_available():92 torch.cuda.manual_seed(seed)93 torch.cuda.manual_seed_all(seed)94 95 96def fed_avg_aggregator(net_list, net_freq):97 sum_parameters = None98 sum_ema_parameters = None99 100 for c in range(len(net_list)):101 global_parameters, global_ema_parameters = net_list[c]102 global_parameters = global_parameters.state_dict()103 global_ema_parameters = global_ema_parameters.state_dict()104 105 if sum_parameters is None:106 sum_parameters = {key: var.clone() * net_freq[c] for key, var in global_parameters.items()}107 else:108 for key in sum_parameters:109 sum_parameters[key] = sum_parameters[key] + global_parameters[key] * net_freq[c]110 111 if sum_ema_parameters is None:112 sum_ema_parameters = {key: var.clone() * net_freq[c] for key, var in global_ema_parameters.items()}113 else:114 for key in sum_ema_parameters:115 sum_ema_parameters[key] = sum_ema_parameters[key] + global_ema_parameters[key] * net_freq[c]116 117 return sum_parameters, sum_ema_parameters118 119 120def build_client_triggers(client_num, img_size, num_points=3, grid_size=8, seed=42, square_size=3):121 assert square_size >= 1 and square_size % 2 == 1, 'square_size should be an odd positive integer, e.g. 3'122 123 margin = max(2, img_size // 8)124 xs = torch.linspace(margin, img_size - margin - 1, steps=grid_size).long()125 ys = torch.linspace(margin, img_size - margin - 1, steps=grid_size).long()126 candidate_positions = [(int(y), int(x)) for y in ys for x in xs]127 128 half = square_size // 2129 triggers = []130 131 for client_id in range(client_num):132 trig = torch.full((3, img_size, img_size), -1.0)133 local_rng = random.Random(seed + client_id)134 selected = local_rng.sample(candidate_positions, num_points)135 136 for (y, x) in selected:137 y1 = max(0, y - half)138 y2 = min(img_size, y + half + 1)139 x1 = max(0, x - half)140 x2 = min(img_size, x + half + 1)141 trig[:, y1:y2, x1:x2] = 1.0142 143 triggers.append(trig)144 145 return triggers146 147 148def maybe_load_global_ckpt():149 if FLAGS.resume_ckpt is None:150 return None, 0151 152 ckpt = torch.load(FLAGS.resume_ckpt, map_location='cpu')153 start_round = int(ckpt.get('round', 0))154 155 print(f'Resume from ckpt: {FLAGS.resume_ckpt}')156 print(f'Start round: {start_round + 1}')157 return ckpt, start_round158 159 160def resolve_logdir(dataset_name: str) -> str:161 if FLAGS.run_name:162 run_name = FLAGS.run_name163 else:164 timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')165 run_name = f'cond{FLAGS.cond_dim}_neg{FLAGS.neg_ratio}_seed{FLAGS.data_distribution_seed}_{timestamp}'166 167 if dataset_name == 'cifar10':168 return os.path.join(FLAGS.base_logdir, 'cifar10_cond_acsctrl', run_name)169 if dataset_name == 'celeba':170 return os.path.join(FLAGS.base_logdir, 'celeba_cond_acsctrl', run_name)171 raise NotImplementedError(f'Unsupported dataset: {dataset_name}')172 173def make_new_logdir(base_logdir):174 timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')175 return os.path.join(176 base_logdir,177 f'cond{FLAGS.cond_dim}_neg{FLAGS.neg_ratio}_seed{FLAGS.data_distribution_seed}_{timestamp}'178 )179 180 181def load_trigger_encoder_from_ckpt(trigger_encoder, global_ckpt):182 if global_ckpt is None:183 return184 if 'trigger_encoder' in global_ckpt:185 trigger_encoder.load_state_dict(global_ckpt['trigger_encoder'], strict=True)186 print('Loaded trigger_encoder from checkpoint.')187 else:188 print('Warning: no trigger_encoder found in checkpoint. Using fresh init.')189 190 191def resolve_client_triggers(global_ckpt):192 if global_ckpt is not None and 'client_triggers' in global_ckpt:193 client_triggers = [t.clone().cpu() for t in global_ckpt['client_triggers']]194 print('Loaded client_triggers from checkpoint.')195 return client_triggers196 197 print('No client_triggers found in checkpoint, building new client_triggers.')198 return build_client_triggers(199 FLAGS.client_num,200 FLAGS.img_size,201 num_points=2,202 grid_size=8,203 seed=42,204 square_size=3,205 )206 207 208def save_client_trigger_images(client_triggers, client_triggers_dir):209 for i, trig in enumerate(client_triggers):210 trig_vis = (trig + 1) / 2211 save_image(trig_vis, os.path.join(client_triggers_dir, f'client_{i}.png'))212 213 trigger_grid = make_grid([(t + 1) / 2 for t in client_triggers], nrow=min(len(client_triggers), 5))214 save_image(trigger_grid, os.path.join(client_triggers_dir, 'all_clients.png'))215 216 217def dump_resume_info(logdir, global_ckpt):218 resume_info_path = os.path.join(logdir, 'resume_info.json')219 payload = {220 'resume_ckpt': FLAGS.resume_ckpt,221 'resume_round': int(global_ckpt.get('round', 0)) if global_ckpt is not None else 0,222 'resume_has_trigger_encoder': bool(global_ckpt is not None and 'trigger_encoder' in global_ckpt),223 'resume_has_client_triggers': bool(global_ckpt is not None and 'client_triggers' in global_ckpt),224 'resume_has_client_states': bool(global_ckpt is not None and 'client_states' in global_ckpt),225 }226 with open(resume_info_path, 'w', encoding='utf-8') as f:227 json.dump(payload, f, indent=2, ensure_ascii=False)228 229 230def train():231 set_seed(FLAGS.seed)232 device = get_device()233 234 if FLAGS.dataset_name == 'cifar10':235 FLAGS.img_size = 32236 elif FLAGS.dataset_name == 'celeba':237 FLAGS.img_size = 64238 else:239 raise NotImplementedError(f'Unsupported dataset: {FLAGS.dataset_name}')240 241 global_ckpt, start_round = maybe_load_global_ckpt()242 FLAGS.logdir = make_new_logdir(FLAGS.base_logdir)243 244 sample_uncond_dir = os.path.join(FLAGS.logdir, 'sample_uncond') 245 sample_correct_dir = os.path.join(FLAGS.logdir, 'sample_correct')246 sample_none_dir = os.path.join(FLAGS.logdir, 'sample_none')247 sample_unseen_dir = os.path.join(FLAGS.logdir, 'sample_unseen')248 ckpt_dir = os.path.join(FLAGS.logdir, 'checkpoints')249 client_triggers_dir = os.path.join(FLAGS.logdir, 'client_triggers')250 train_metrics_csv = os.path.join(FLAGS.logdir, 'train_metrics.csv')251 sample_metrics_csv = os.path.join(FLAGS.logdir, 'sample_metrics.csv')252 client_data_stats_csv = os.path.join(FLAGS.logdir, 'client_data_stats.csv')253 254 os.makedirs(FLAGS.logdir, exist_ok=True)255 os.makedirs(ckpt_dir, exist_ok=True)256 257 if FLAGS.mode != 'unconditional':258 os.makedirs(client_triggers_dir, exist_ok=True)259 260 if FLAGS.save_samples and FLAGS.sample_every > 0:261 if FLAGS.mode == 'unconditional':262 os.makedirs(sample_uncond_dir, exist_ok=True)263 else:264 os.makedirs(sample_correct_dir, exist_ok=True)265 os.makedirs(sample_none_dir, exist_ok=True)266 os.makedirs(sample_unseen_dir, exist_ok=True)267 268 print('device:', device)269 print('mode:', FLAGS.mode)270 print('dataset:', FLAGS.dataset_name)271 print('cond_dim:', FLAGS.cond_dim)272 print('neg_ratio:', FLAGS.neg_ratio)273 print('logdir:', FLAGS.logdir)274 275 model_cond_dim = None if FLAGS.mode == 'unconditional' else FLAGS.cond_dim276 277 net_model_global = UNet(278 T=FLAGS.T,279 ch=FLAGS.ch,280 ch_mult=FLAGS.ch_mult,281 attn=FLAGS.attn,282 num_res_blocks=FLAGS.num_res_blocks,283 dropout=FLAGS.dropout,284 cond_dim=model_cond_dim,285 ).to(device)286 287 if FLAGS.mode == 'unconditional':288 trigger_encoder = None289 client_triggers = None290 else:291 trigger_encoder = TriggerEncoder(292 in_ch=3,293 base_ch=FLAGS.trigger_base_ch,294 cond_dim=FLAGS.cond_dim,295 img_size=FLAGS.img_size,296 grid_size=8,297 hidden_dim=256,298 use_coord_feat=True,299 ).to(device)300 301 load_trigger_encoder_from_ckpt(trigger_encoder, global_ckpt)302 client_triggers = resolve_client_triggers(global_ckpt)303 save_client_trigger_images(client_triggers, client_triggers_dir)304 305 x_T = torch.randn(FLAGS.sample_batch_size, 3, FLAGS.img_size, FLAGS.img_size, device=device)306 307 writer = SummaryWriter(FLAGS.logdir) if FLAGS.save_tensorboard else None308 if writer is not None:309 writer.flush()310 311 with open(os.path.join(FLAGS.logdir, 'flagfile.txt'), 'w', encoding='utf-8') as f:312 f.write(FLAGS.flags_into_string())313 314 dump_resume_info(FLAGS.logdir, global_ckpt)315 316 with open(train_metrics_csv, 'w', newline='') as f:317 writer_csv = csv.writer(f)318 writer_csv.writerow([319 'round', 'client_id', 'avg_train_loss', 'last_train_loss', 'lr',320 'step_count', 'num_batches', 'num_samples'321 ])322 323 with open(sample_metrics_csv, 'w', newline='') as f:324 writer_csv = csv.writer(f)325 writer_csv.writerow([326 'round', 'client_id', 'trigger_mode', 'sample_mean', 'sample_std',327 'sample_min', 'sample_max'328 ])329 330 clients_group = ClientsGroupMultiTargetAttackedNonIID(331 dataset_name=FLAGS.dataset_name,332 batch_size=FLAGS.batch_size,333 clients_num=FLAGS.client_num,334 device=device,335 mode=FLAGS.mode,336 trigger_encoder=trigger_encoder,337 client_triggers=client_triggers,338 neg_ratio=FLAGS.neg_ratio,339 cond_dim=FLAGS.cond_dim,340 split_mode=FLAGS.split_mode,341 unseen_num_points=FLAGS.unseen_num_points,342 unseen_grid_size=FLAGS.unseen_grid_size,343 unseen_square_size=FLAGS.unseen_square_size,344 unseen_base_seed=FLAGS.unseen_base_seed,345 blur_kernel_size=FLAGS.blur_kernel_size,346 blur_repeat=FLAGS.blur_repeat,347 data_distribution_seed=FLAGS.data_distribution_seed,348 ema_scale=FLAGS.ema_scale,349 num_workers=FLAGS.num_workers,350 beta_1=FLAGS.beta_1,351 beta_T=FLAGS.beta_T,352 T=FLAGS.T,353 )354 355 for i in range(FLAGS.client_num):356 clients_group.clients_set[i].init(357 net_model_global,358 FLAGS.lr,359 FLAGS.parallel,360 global_ckpt=global_ckpt,361 )362 363 client_idx = [i for i in range(FLAGS.client_num)]364 clients_targets = [clients_group.clients_set[c].get_targets_num() for c in client_idx]365 train_data_sum = sum(clients_targets)366 367 with open(client_data_stats_csv, 'w', newline='') as f:368 writer_csv = csv.writer(f)369 writer_csv.writerow(['client_id', 'num_samples'])370 for c in client_idx:371 writer_csv.writerow([c, clients_targets[c]])372 373 for round_idx in range(start_round, FLAGS.total_round):374 tqdm.write(f'\n[Round {round_idx + 1}/{FLAGS.total_round}]')375 net_freq = [clients_targets[c] / train_data_sum for c in client_idx]376 377 net_list = []378 round_train_metrics = []379 380 for c in client_idx:381 model_state, ema_state, train_metrics = clients_group.clients_set[c].local_train(382 round_idx,383 FLAGS.local_epoch,384 mid_T=FLAGS.mid_T,385 )386 net_list.append((model_state, ema_state))387 round_train_metrics.append(train_metrics)388 389 with open(train_metrics_csv, 'a', newline='') as f:390 writer_csv = csv.writer(f)391 for m in round_train_metrics:392 writer_csv.writerow([393 m['round'], m['client_id'], m['avg_train_loss'], m['last_train_loss'],394 m['lr'], m['step_count'], m['num_batches'], m['num_samples']395 ])396 397 sum_parameters, sum_ema_parameters = fed_avg_aggregator(net_list, net_freq)398 399 for c in client_idx:400 clients_group.clients_set[c].set_global_parameters(sum_parameters, sum_ema_parameters)401 402 should_sample = FLAGS.save_samples and FLAGS.sample_every > 0 and ((round_idx + 1) % FLAGS.sample_every == 0)403 round_sample_metrics = []404 405 if should_sample:406 if FLAGS.mode == 'unconditional':407 samples_uncond = []408 409 for c in client_idx:410 client = clients_group.clients_set[c]411 with torch.no_grad():412 x_uncond = client.get_sample(413 x_T,414 FLAGS.sample_start_step,415 FLAGS.sample_end_step,416 trigger_mode='unconditional',417 )418 419 samples_uncond.append(x_uncond)420 421 round_sample_metrics.append({422 'round': round_idx + 1,423 'client_id': c,424 'trigger_mode': 'unconditional',425 'sample_mean': float(x_uncond.mean().item()),426 'sample_std': float(x_uncond.std().item()),427 'sample_min': float(x_uncond.min().item()),428 'sample_max': float(x_uncond.max().item()),429 })430 431 samples_uncond = torch.cat(samples_uncond, dim=0)432 grid_uncond = (make_grid(samples_uncond, nrow=x_T.size(0)) + 1) / 2433 434 save_image(grid_uncond, os.path.join(sample_uncond_dir, f'{round_idx + 1}.png'))435 436 if writer is not None:437 writer.add_image('sample_uncond', grid_uncond, round_idx + 1)438 439 else:440 samples_correct = []441 samples_none = []442 samples_unseen = []443 444 for c in client_idx:445 client = clients_group.clients_set[c]446 with torch.no_grad():447 x_correct = client.get_sample(448 x_T, FLAGS.sample_start_step, FLAGS.sample_end_step, trigger_mode='correct'449 )450 x_none = client.get_sample(451 x_T, FLAGS.sample_start_step, FLAGS.sample_end_step, trigger_mode='none'452 )453 x_unseen = client.get_sample(454 x_T, FLAGS.sample_start_step, FLAGS.sample_end_step, trigger_mode='unseen', unseen_seed=999455 )456 samples_correct.append(x_correct)457 samples_none.append(x_none)458 samples_unseen.append(x_unseen)459 460 for mode_name, sample in [('correct', x_correct), ('none', x_none), ('unseen', x_unseen)]:461 round_sample_metrics.append({462 'round': round_idx + 1,463 'client_id': c,464 'trigger_mode': mode_name,465 'sample_mean': float(sample.mean().item()),466 'sample_std': float(sample.std().item()),467 'sample_min': float(sample.min().item()),468 'sample_max': float(sample.max().item()),469 })470 471 samples_correct = torch.cat(samples_correct, dim=0)472 samples_none = torch.cat(samples_none, dim=0)473 samples_unseen = torch.cat(samples_unseen, dim=0)474 475 grid_correct = (make_grid(samples_correct, nrow=x_T.size(0)) + 1) / 2476 grid_none = (make_grid(samples_none, nrow=x_T.size(0)) + 1) / 2477 grid_unseen = (make_grid(samples_unseen, nrow=x_T.size(0)) + 1) / 2478 479 save_image(grid_correct, os.path.join(sample_correct_dir, f'{round_idx + 1}.png'))480 save_image(grid_none, os.path.join(sample_none_dir, f'{round_idx + 1}.png'))481 save_image(grid_unseen, os.path.join(sample_unseen_dir, f'{round_idx + 1}.png'))482 483 if writer is not None:484 writer.add_image('sample_correct', grid_correct, round_idx + 1)485 writer.add_image('sample_none', grid_none, round_idx + 1)486 writer.add_image('sample_unseen', grid_unseen, round_idx + 1)487 488 with open(sample_metrics_csv, 'a', newline='') as f:489 writer_csv = csv.writer(f)490 for m in round_sample_metrics:491 writer_csv.writerow([492 m['round'], m['client_id'], m['trigger_mode'], m['sample_mean'],493 m['sample_std'], m['sample_min'], m['sample_max']494 ])495 496 avg_round_loss = sum(m['avg_train_loss'] for m in round_train_metrics) / max(1, len(round_train_metrics))497 print(f'[Round {round_idx + 1}] avg_train_loss={avg_round_loss:.6f}')498 499 if FLAGS.save_round > 0 and (round_idx + 1) % FLAGS.save_round == 0:500 client_states = None501 if FLAGS.save_client_states:502 client_states = {}503 for c in client_idx:504 client_states[f'client_{c}'] = clients_group.clients_set[c].export_client_state()505 506 max_step_count = max(clients_group.clients_set[c]._step_count for c in client_idx)507 508 global_ckpt_to_save = {509 'round': round_idx + 1,510 'step_count': max_step_count,511 'config': {512 'mode': FLAGS.mode,513 'dataset_name': FLAGS.dataset_name,514 'img_size': FLAGS.img_size,515 'cond_dim': None if FLAGS.mode == 'unconditional' else FLAGS.cond_dim,516 'client_num': FLAGS.client_num,517 'neg_ratio': FLAGS.neg_ratio,518 'split_mode': FLAGS.split_mode,519 'data_distribution_seed': FLAGS.data_distribution_seed,520 'T': FLAGS.T,521 'ch': FLAGS.ch,522 'ch_mult': list(FLAGS.ch_mult),523 'attn': list(FLAGS.attn),524 'num_res_blocks': FLAGS.num_res_blocks,525 'dropout': FLAGS.dropout,526 'lr': FLAGS.lr,527 'mid_T': FLAGS.mid_T,528 'local_epoch': FLAGS.local_epoch,529 'save_round': FLAGS.save_round,530 'resumed_from': FLAGS.resume_ckpt,531 },532 }533 534 if FLAGS.mode != 'unconditional':535 global_ckpt_to_save['trigger_encoder'] = trigger_encoder.state_dict()536 global_ckpt_to_save['client_triggers'] = [t.cpu() for t in client_triggers]537 538 global_ckpt_to_save['config'].update({539 'trigger_base_ch': FLAGS.trigger_base_ch,540 'trigger_img_size': FLAGS.img_size,541 'trigger_grid_size': 8,542 'trigger_hidden_dim': 256,543 'trigger_use_coord_feat': True,544 'unseen_num_points': FLAGS.unseen_num_points,545 'unseen_grid_size': FLAGS.unseen_grid_size,546 'unseen_square_size': FLAGS.unseen_square_size,547 'unseen_base_seed': FLAGS.unseen_base_seed,548 'blur_kernel_size': FLAGS.blur_kernel_size,549 'blur_repeat': FLAGS.blur_repeat,550 })551 552 if FLAGS.save_global_model:553 global_ckpt_to_save['global_model'] = sum_parameters554 if FLAGS.save_global_ema_model:555 global_ckpt_to_save['global_ema_model'] = sum_ema_parameters556 if client_states is not None:557 global_ckpt_to_save['client_states'] = client_states558 559 ckpt_path = os.path.join(ckpt_dir, f'global_ckpt_round{round_idx + 1}.pt')560 torch.save(global_ckpt_to_save, ckpt_path)561 print(f'Saved checkpoint: {ckpt_path}')562 563 if writer is not None:564 writer.close()565 566 567def main(argv):568 del argv569 warnings.simplefilter(action='ignore', category=FutureWarning)570 if FLAGS.train:571 train()572 else:573 print('Use --train to start training.')574 print('Smoke test example: python train_min.py --train --smoke_test --dataset_name=cifar10 --device=cuda')575 576 577if __name__ == '__main__':578 app.run(main)579