(DIST=True, graph=True, sequential=False, verbosity=0)
| 31 | |
| 32 | |
| 33 | def train_resnet(DIST=True, graph=True, sequential=False, verbosity=0): |
| 34 | |
| 35 | # Define the hypermeters for the train_resnet |
| 36 | niters = 100 |
| 37 | batch_size = 32 |
| 38 | sgd = opt.SGD(lr=0.1, momentum=0.9, weight_decay=1e-5) |
| 39 | |
| 40 | IMG_SIZE = 224 |
| 41 | |
| 42 | # For distributed training, sequential has better throughput in the current version |
| 43 | if DIST == True: |
| 44 | sgd = opt.DistOpt(sgd) |
| 45 | world_size = sgd.world_size |
| 46 | local_rank = sgd.local_rank |
| 47 | global_rank = sgd.global_rank |
| 48 | sequential = True |
| 49 | else: |
| 50 | local_rank = 0 |
| 51 | world_size = 1 |
| 52 | global_rank = 0 |
| 53 | sequential = False |
| 54 | |
| 55 | dev = device.create_cuda_gpu_on(local_rank) |
| 56 | |
| 57 | tx = tensor.Tensor((batch_size, 3, IMG_SIZE, IMG_SIZE), dev) |
| 58 | ty = tensor.Tensor((batch_size,), dev, tensor.int32) |
| 59 | x = np.random.randn(batch_size, 3, IMG_SIZE, IMG_SIZE).astype(np.float32) |
| 60 | y = np.random.randint(0, 1000, batch_size, dtype=np.int32) |
| 61 | tx.copy_from_numpy(x) |
| 62 | ty.copy_from_numpy(y) |
| 63 | |
| 64 | dev.SetVerbosity(verbosity) |
| 65 | dev.SetSkipIteration(5) |
| 66 | |
| 67 | # Construct the model |
| 68 | from model import resnet |
| 69 | model = resnet.resnet50(num_channels=3, num_classes=1000) |
| 70 | |
| 71 | model.train() |
| 72 | model.set_optimizer(sgd) |
| 73 | model.compile([tx], is_train=True, use_graph=graph, sequential=sequential) |
| 74 | |
| 75 | # Train model |
| 76 | dev.Sync() |
| 77 | start = time.time() |
| 78 | with trange(niters) as t: |
| 79 | for _ in t: |
| 80 | model(tx, ty, dist_option='fp32', spars=None) |
| 81 | |
| 82 | dev.Sync() |
| 83 | end = time.time() |
| 84 | titer = (end - start) / float(niters) |
| 85 | throughput = float(niters * batch_size * world_size) / (end - start) |
| 86 | if global_rank == 0: |
| 87 | print("\nThroughput = {} per second".format(throughput), flush=True) |
| 88 | print("TotalTime={}".format(end - start), flush=True) |
| 89 | print("Total={}".format(titer), flush=True) |
| 90 | dev.PrintTimeProfiling() |
no test coverage detected